/*  Program:     tcpnat.cpp
    Author:      Rob Beverly <rbeverly \at\ mit \d0t\ edu>
    Date:        February 9, 2004
    Updated:     
    Description: NAT-ify a tcpdump trace.  From input trace, we 
                 remap sources such that output sources = 
                 input sources / natFac.  Simulates natFac:1 NAT.
*/

#include <stdio.h>
#include <unistd.h>

#define __FAVOR_BSD
#include <pcap.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/in_systm.h>
#include <netinet/ip.h>
#include <arpa/inet.h>

#include <iostream>
#include <map>

#define VERSION "0.1"

using namespace std;

class in_addrLess {
    public:
        bool operator() (const in_addr i1, const in_addr i2) const {
            if (i1.s_addr > i2.s_addr) return true;
            return false;
        }
};

typedef map<struct in_addr, struct in_addr, in_addrLess> addressMap;
typedef addressMap::iterator ipMapIter;


void usage(char *prog) {
    printf("usage: %s [-v] [-f<natfac>] -i<input> -o<output>\n", prog);
    exit(-1);
}


int headerLen(int type) {
  switch(type){
    case DLT_NULL:
      return 4;
      break;
    case DLT_RAW:
      return 0;
      break;
    case DLT_EN10MB:
      return 14;
    default:
      printf("Unknown data-link type: %d\n", type);
      return 0;
      break;
  }
}


int main (int argc, char **argv) {
    pcap_t *pcap;
    pcap_dumper_t *pd;
    struct pcap_pkthdr hdr;
    const u_char *pkt = NULL;
    struct ip *ippkt = NULL;
    char errbuf[256];
    char *input = NULL;
    char *output = NULL;
    int ch;

    static struct in_addr lastNew;
    int uniqueIPs = 0;
    int natFac = 2;
    int pcapHL = 0;
    bool verbose = false;
    addressMap ipMap;

    if (argc < 2) usage(argv[0]);

    cout << argv[0] << " v" << VERSION " - NATify tcpdump trace.\n";

    /* Parse the command-line. */
    while ((ch = getopt(argc, argv, "hf:i:o:v")) != EOF) {
        switch ((char) ch) {
        case 'h':
            usage(argv[0]);
            break;
        case 'i':
            input = optarg;
            break;
        case 'o':
            output = optarg;
            break;
        case 'f':
            natFac = atoi(optarg);
            break;
        case 'v':
            verbose = true;
            break;
        }
    }
 
    if (!input || !output) usage(argv[0]);
    cout << "Using NAT factor of: " << natFac << endl;

    /* Open pcap files */
    cout << "Opening input: (" << input << ").\n";
    if ((pcap = pcap_open_offline(input, errbuf)) == NULL) {
        printf("Error: %s\n", errbuf);
        exit(-1);
    }
    cout << "Opening output: (" << output << ").\n";
    if ((pd = pcap_dump_open(pcap, output)) == NULL) {
        exit(-1);
    }

    pcapHL = headerLen(pcap_datalink(pcap));
    while ( (pkt = pcap_next(pcap, &hdr)) ) {
        ippkt = (struct ip *) (pkt + pcapHL);
        if (ipMap.find(ippkt->ip_src) == ipMap.end()) {
            if (uniqueIPs % natFac == 0) {
                lastNew=(ippkt->ip_src);
            }
            uniqueIPs++;
            ipMap[(ippkt->ip_src)] = lastNew;
            if (verbose) {
                cout << "Adding: " << inet_ntoa(ippkt->ip_src) << " (" << uniqueIPs;
                cout << ") -> " << inet_ntoa(lastNew) << endl;
            }
        }
        ippkt->ip_src = ipMap[(ippkt->ip_src)];
        pcap_dump((u_char *)pd, &hdr, pkt);
    } 
    pcap_close(pcap);
    pcap_dump_close(pd);

    if (verbose) {
        for (ipMapIter iter=ipMap.begin(); iter != ipMap.end(); iter++) {
            struct in_addr a = iter->first;
            struct in_addr b = iter->second;
            cout << "IP: " << inet_ntoa(a);
            cout << " Maps to: " << inet_ntoa(b) << endl;
        }
    }

    cout << "Done.  Mapped " << ipMap.size() << " sources." << endl;
    exit(0);
}