/*
 * Copyright (C) 2015 by the University of Southern California
 * $Id: dns_util.c 23977 2015-05-13 17:56:01Z 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 <err.h>
#include <stdio.h>
#include <ldns/ldns.h>
#include "dns_util.h"
#include "global_var.h"

#include <netinet/ip.h>
#include <netinet/udp.h>
#include <netinet/tcp.h>

int 
format_dns_query(uint8_t **pkt, size_t *pkt_size, query_t *q)
{ 
  return format_dns_query(pkt, pkt_size, q, 0);
}

int 
format_dns_query(uint8_t **pkt, size_t *pkt_size, query_t *q, uint16_t z)
{
  const char    *func_name = "format_dns_query";
  ldns_pkt      *p;
  ldns_rr_type  rr_type    = ldns_get_rr_type_by_name(q->qtype);
  ldns_rr_class rr_class   = ldns_get_rr_class_by_name(q->qclass);
  ldns_status   s;
  s = ldns_pkt_query_new_frm_str(&p, q->qname, rr_type, rr_class , (uint16_t)(LDNS_RD|LDNS_CD));
  if(s != LDNS_STATUS_OK) {
    fprintf (stderr, "%s: [error] err by ldns(%s)", func_name, ldns_get_errorstr_by_id(s));
    return 1;
  }
  if(!p) {
    fprintf (stderr, "%s: [error] out of memory", func_name);
    return 1;
  }
  ldns_pkt_set_random_id(p);
  ldns_pkt_set_edns_z(p, z);
  s = ldns_pkt2wire(pkt, p, pkt_size);
  if(s != LDNS_STATUS_OK) {
    fprintf (stderr, "%s: [error] %s",func_name, ldns_get_errorstr_by_id(s));
    return 1;
  }
  ldns_pkt_free(p);
  return 0;
}

int
create_dns_raw(hdr_t *h, int conn, uint8_t *buf, int buf_sz, query_t *q)
{
  //printf("conn=%d\n", conn);
  const char    *func_name = "create_dns_raw";
  memset(buf, 0, buf_sz);
  //int sz             = 0;
  int l              = 0;
  struct ip *ip      = NULL;
  struct udphdr *udp = NULL;
  uint8_t *dns       = NULL;
  size_t dns_sz      = 0;

  ip = (struct ip *) &buf[l];
  l += 20;
  ip->ip_v = 4;
  ip->ip_hl = 5;
  ip->ip_tos = 0;
  ip->ip_id = random() & 0xffff;
  ip->ip_off = 0;
  ip->ip_ttl = 64;
  ip->ip_p = (conn==CONN_TCP ? IPPROTO_TCP : IPPROTO_UDP);
  ip->ip_sum = 0;
  ip->ip_src = h->ip_src;
  ip->ip_dst = h->ip_dst;

  udp = (struct udphdr *) &buf[l];
  l += 8;
  udp->source = h->spt;
  udp->dest = h->dpt;

  if (format_dns_query(&dns, &dns_sz, q) != 0)
    errx(1, "%s: [error] format_dns_query fails", func_name);
  memcpy (&buf[l], dns, dns_sz);
  free(dns);
  
  l+= dns_sz;
  udp->len = htons(l - 20);
  ip->ip_len = htons(l);
  //ip->ip_sum = cksum(buf,l);
  udp->check = udp_tcp_cksum(&buf[20], l-20, h->ip_src, h->ip_dst, ip->ip_p);
  return l;
}
int
create_dns_tcp(uint8_t *buf, int buf_sz, query_t *q)
{
  memset(buf, 0, buf_sz);
  uint8_t *dns       = NULL;
  size_t dns_sz      = 0;
  if (format_dns_query(&dns, &dns_sz, q) != 0)
    errx(1, "format_dns_query fails");
  uint16_t len_to_write = htons(dns_sz);
  memcpy(buf, &len_to_write, sizeof(len_to_write));
  memcpy (&buf[sizeof(len_to_write)], dns, dns_sz);
  free(dns);
  return dns_sz + 2;
}

void
print_dns_pkt(uint8_t *buf, int buf_sz)
{
  ldns_status status;
  ldns_pkt* pkt = NULL;
  //parse the reply packet
  status = ldns_wire2pkt(&pkt, buf, (size_t)buf_sz);
  if(status != LDNS_STATUS_OK) {
    ldns_get_errorstr_by_id(status);
    errx(1, "cannot parse the reply packet");
  }
  //print the pkt
  ldns_pkt_print(stdout, pkt);
  ldns_pkt_free(pkt);
}

bool
has_to_bit(uint8_t *p, size_t l)
{
  ldns_status status;
  ldns_pkt* pkt = NULL;
  uint8_t *msg = new uint8_t[l];
  memcpy(msg, p, l);
  status = ldns_wire2pkt(&pkt, msg, l);
  if(status != LDNS_STATUS_OK) {
    ldns_get_errorstr_by_id(status);
    warnx("cannot parse the packet");
    return false;
  }
  assert(pkt != NULL);
  if(!ldns_pkt_edns(pkt)) {
    DBG("NO edns in packet\n");
    ldns_pkt_free(pkt);
    return false;
  }
  bool found = false;
  ldns_rdf* qname_rdf = NULL;
  char* qname_str = NULL;
  ldns_rr *q_rr = NULL;
  for (int i = 0; i < ldns_pkt_qdcount(pkt); i++) {
    q_rr = ldns_rr_list_rr(ldns_pkt_question(pkt), i);
    qname_rdf = ldns_rr_owner(q_rr);
    qname_str = ldns_rdf2str(qname_rdf);
    DBG("qname=[%s]\n", qname_str);
    if (strcmp(qname_str, "STARTTLS") == 0 ||
	strcmp(qname_str, "STARTTLS.") == 0)
      found = true;
    if (qname_str) {
      LDNS_FREE(qname_str);
      qname_str = NULL;
    }
    if (found)
      break;
  }
  uint16_t zbit = ldns_pkt_edns_z(pkt); //get the z bit
  ldns_pkt_free(pkt);
  delete[] msg;
  return found && (zbit & EDNS_Z_TO);
}

//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;
}

//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;
}

