/*
 * Copyright (C) 2015 by the University of Southern California
 * $Id: socket_util.c 24224 2015-05-28 22:17:46Z liangzhu $
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License,
 * version 2, as published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the 
 * GNU General Public License for more details.                                      
 *
 * You should have received a copy of the GNU General Public License along           
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA. 
 * 
 */

#include <unistd.h>
#include <stdio.h>
#include <sys/socket.h>
//#include <netinet/ip.h>
//#include <netinet/udp.h>
//#include <netinet/tcp.h>
#include <err.h>
#include <string.h>
#include "socket_util.h"
#include <arpa/inet.h>

//#define MAX_PKT_LEN 4096

//This is copied from https://tools.ietf.org/html/rfc1071
uint16_t 
cksum (uint16_t *addr, int count)
{
  uint16_t checksum = 0;
  /* Compute Internet Checksum for "count" bytes
   *         beginning at location "addr".
   */
  register long sum = 0;
  
  while( count > 1 )  {
    /*  This is the inner loop */
    sum += *addr++;
    count -= 2;
  }
  
  /*  Add left-over byte, if any */
  if( count > 0 )
    sum += *addr;
  
  /*  Fold 32-bit sum to 16 bits */
  while (sum>>16)
    sum = (sum & 0xffff) + (sum >> 16);

  checksum = ~sum;
  return checksum;
}

//create a raw socket, udp or tcp
int
create_raw_socket(int one)
{
  int raw = -1;
  if ((raw = socket(AF_INET, SOCK_RAW, IPPROTO_RAW)) < 0)
    err(1, "SOCK_RAW");
  if (setsockopt(raw, IPPROTO_IP, IP_HDRINCL, &one, sizeof(one)) < 0)
    err(1, "setsockopt IP_HDRINCL");
  return raw;
}

//udp_cksum
uint16_t
udp_tcp_cksum (uint8_t *pkt, uint16_t len, struct in_addr ip_src, struct in_addr ip_dst, uint8_t protocol)
{
  pseudo_hdr fake_hdr;
  memset(&fake_hdr, 0, sizeof(pseudo_hdr));
  fake_hdr.ip_src.s_addr = ip_src.s_addr;
  fake_hdr.ip_dst.s_addr = ip_dst.s_addr;
  fake_hdr.reserved = 0;
  fake_hdr.protocol = protocol;
  fake_hdr.length = htons(len);
  int padding = 0;
  if (len % 2 != 0)
    padding = 1;
  int buf_len = len + sizeof(pseudo_hdr) + padding;
  uint8_t *buf = new uint8_t[buf_len];
  memset(buf, 0, buf_len);
  memcpy(buf, &fake_hdr, sizeof(pseudo_hdr));
  memcpy(buf + sizeof(pseudo_hdr), pkt, len);
  uint16_t sum = cksum((uint16_t*)buf, buf_len);
  delete[] buf;
  return sum;
}

//make a connection either for udp or tcp
int mk_conn (const char *ip, int port, bool tcp)
{
  struct sockaddr_in sa;
  int fd = -1;
  if ((fd = socket(AF_INET, tcp ? SOCK_STREAM : SOCK_DGRAM, 0)) == -1)
    err(1, "socket");
  memset (&sa, '\0', sizeof (sa));
  sa.sin_family = AF_INET;
  sa.sin_port = htons (port);

  //if (strcmp("ANY", ip) || strcmp("any", ip))
  //  sa.sin_addr.s_addr = htonl(INADDR_ANY);
  //else
  inet_pton (AF_INET, ip, &sa.sin_addr);
  if (connect (fd, (struct sockaddr *) & sa, sizeof (sa)) < 0)
    err(1, "connect");
  return fd;
}

//is const char * a ipv4 address?
bool ipv4_addr(const char *s)
{
  struct sockaddr_in sa;
  return inet_pton(AF_INET, s, &(sa.sin_addr)) != 0;
}

//is const char * a ipv6 address?
bool ipv6_addr(const char *s)
{
  struct sockaddr_in6 sa;
  return inet_pton(AF_INET6, s, &(sa.sin6_addr)) != 0;
}

char *ip_b2s(const struct sockaddr *sa, char *s, size_t maxlen)
{
  const char *rv = NULL;
  switch(sa->sa_family) {
  case AF_INET:
    rv = inet_ntop(AF_INET, &(((struct sockaddr_in *)sa)->sin_addr), s, maxlen);
    if (rv == NULL)
      err(1, "[error] inet_ntop(AF_INET) fails, abort!");
    break;
  case AF_INET6:
    rv = inet_ntop(AF_INET6, &(((struct sockaddr_in6 *)sa)->sin6_addr), s, maxlen);
    if (rv == NULL)
     err(1, "[error] inet_ntop(AF_INET6) fails, abort!");
    break;
  default:
    errx(1, "[error] AF is unknown!");
    return NULL;
  }
  return s;
}

int ip_s2b(const char *ip_v, const char *s, const struct sockaddr *sa)
{
  int rv = -1;
  if (ip_v == NULL) {
    errx(1, "input for ip_s2b is NULL");
    return -1; 
  }
  if (strcmp(ip_v, "ipv4") == 0)
    rv = inet_pton(AF_INET, s, &(((struct sockaddr_in *)sa)->sin_addr));
  else if (strcmp(ip_v, "ipv6") == 0)
    rv = inet_pton(AF_INET6, s, &(((struct sockaddr_in6 *)sa)->sin6_addr));
  else
    errx(1, "ip version [%s] is unknow in ip_s2b", ip_v);
  if (rv != 1)
    errx(1, "ip_s2b [%s] fails", s);
  return rv;
}

void clean_fd(int f)
{
  if (f > 0)
    close(f);
}

int udp_dns_pkt(int s, uint8_t *buf, int buf_len, struct sockaddr* sa, size_t* sa_len)
{
  memset(buf, 0, buf_len);
  ssize_t recv_len = recvfrom(s, buf, buf_len, 0, sa, (socklen_t*)sa_len);
  if (recv_len == -1)
    err(1, "[error] recvfrom return -1");
    //perror("recvfrom");
  return recv_len;
}

int pull_dns_pkt(void *s, uint8_t *buf, int buf_len, const char *c)
{
  //check the protocol name
  bool tls = (strcmp(c, "tls") == 0 ||
	      strcmp(c, "ssl") == 0 ||
	      strcmp(c, "TLS") == 0 ||
	      strcmp(c, "SSL") == 0);
  if (!tls &&
      (strcmp(c, "TCP")) &&
      (strcmp(c, "tcp")))
    errx(1, "[error] connection type[%s] is unknown, abort!", c);

  uint16_t len = 0; //DNS payload length given in TCP packet
  int recv_len = -1;

  //read the first two byte for length
  if (!tls)
    recv_len = recv(*(int *)s, (void*)&len, sizeof(len), 0);
  else {
#ifdef USE_GNUTLS
    recv_len = gnutls_record_recv((gnutls_session_t)s, (void*)&len, (int)sizeof(len));
#elseif USE_OPENSSL
    SSL_read((SSL*)s, (void*)&len, (int)sizeof(len));
#endif 
  }

  //connection closed? error?
  if (recv_len == 0)                                                                                               
    return 0;
  if (recv_len < (ssize_t)sizeof(len))
    err(1, "[error] recv[%d]< sizeof(len)[%lu], abort!", recv_len, sizeof(len));

  //read the rest of bytes
  len = ntohs(len);
  memset(buf, 0, buf_len);
  if (!tls)
    recv_len = recv(*(int *)s, buf, len, 0);
  else {
#ifdef USE_GNUTLS
    recv_len = gnutls_record_recv((gnutls_session_t)s, buf, len);
#elseif USE_OPENSSL
    SSL_read((SSL*)s, buf, len);
#endif 
  }

  if (recv_len < len)
    err(1, "[error] recv[%d]<len[%d], abort!", recv_len, len);
  return len;
}
