/* tcpflow.c
 *
 * Tool used to reconstruct dumped tcp packets and extract images
 *
 *
 */

#include <pcap.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <string.h>

// these are for cross-platform compatibility with mac.
// if you keep these defines, use the th_sport-style headers
// rather than the 'source' (linux)-style headers.

#ifndef __USE_MISC
#define __USE_MISC 1
#endif

#ifndef __FAVOR_BSD
#define __FAVOR_BSD 1
#endif

#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>

struct tcpflow
{
    struct in_addr ip_src, ip_dst;
    u_int16_t src_port, dst_port;
    tcp_seq initial_seq;
    unsigned int packet_count, bytes_count;
    int fd;
    char ip_src_s[16];
    char ip_dst_s[16];
    char session_fname[64];
    struct tcpflow *next;
};
typedef struct tcpflow tcpflow_t;

tcpflow_t *head = NULL;
const char *basedir = NULL;

int check_session(tcpflow_t *list, struct in_addr ip_src, struct in_addr ip_dst,
                  u_int16_t src_port, u_int16_t dst_port, tcpflow_t **session_ptr)
{
    tcpflow_t *iter = list; /* iterator */
    while (iter)
    {
        if (iter->ip_src.s_addr == ip_src.s_addr && iter->ip_dst.s_addr == ip_dst.s_addr &&
            iter->src_port == src_port && iter->dst_port == dst_port)
        {
            *session_ptr = iter;
            return 0;
        }
        iter = iter->next;
    }

    return 1;
}

/* this function shall be only called when new connection happens */
int add_session(tcpflow_t **list, struct ip *ip_pkt, struct tcphdr *tcp_pkt)
{
    if (!(tcp_pkt->th_flags & TH_SYN))
    {
        fprintf(stderr, "skipping non-syn packet\n");
        return 0;
    }

    tcpflow_t *node = (tcpflow_t *) malloc(sizeof(tcpflow_t));
    if (node)
    {
        memset(node, 0, sizeof(tcpflow_t));
        /* fill node with data */
        memcpy(&(node->ip_src), &(ip_pkt->ip_src), sizeof(struct in_addr));
        memcpy(&(node->ip_dst), &(ip_pkt->ip_dst), sizeof(struct in_addr));
        node->src_port = ntohs(tcp_pkt->th_sport);
        node->dst_port = ntohs(tcp_pkt->th_dport);
        node->initial_seq = ntohl(tcp_pkt->th_seq);
        node->bytes_count += ((ntohs(ip_pkt->ip_len) - 20 - (tcp_pkt->th_off << 2))); /* calc payload */
        strcpy(node->ip_src_s, inet_ntoa(ip_pkt->ip_src));
        strcpy(node->ip_dst_s, inet_ntoa(ip_pkt->ip_dst));

        char buff[64] = { 0 };
        sprintf(buff, "%s.%u-%s.%u.logs", node->ip_src_s, node->src_port,
        node->ip_dst_s, node->dst_port);
        strncpy(node->session_fname, buff, 64);

        /* open file for writing */
        char fname[300] = { 0 };
        sprintf(fname, "%s%s%s", basedir, basedir[strlen(basedir)-1] == '/' ? "" : "/", node->session_fname);
        if ((node->fd = open(fname, O_WRONLY | O_CREAT, 0644)) == -1)
            perror("failed to create file for tcp session");

        node->next = *list;
        *list = node;
    }
    else
    {
        perror("failed to create session node");
        return -1;
    }

    return 0;
}

int cont_session(tcpflow_t *curr_session, struct ip *ip_pkt, struct tcphdr *tcp_pkt)
{
    /* calc payload */
    unsigned int payload = (ntohs(ip_pkt->ip_len) - 20 - (tcp_pkt->th_off << 2));

    if (!payload && (tcp_pkt->th_flags & TH_ACK)) /* silly to check for ACK flag? It's always set after handshake, eh */
    {
        /* or possibly finishing a 3-way handshake */
        fprintf(stderr, "%s is ACKing\n", curr_session->ip_src_s);
        return 0;
    }

    /* append payload, if any */
    if (payload && curr_session->fd)
    {
        void *payload_addr = ((char *) (tcp_pkt)) + (tcp_pkt->th_off << 2);
        /* seek to some offset from beginning */
        lseek(curr_session->fd, ntohl(tcp_pkt->th_seq) - curr_session->initial_seq, SEEK_SET);
        write(curr_session->fd, payload_addr, payload);
        curr_session->bytes_count += payload;

        if (tcp_pkt->th_flags & TH_PUSH)
        {   /* we should still keep fd open since we might communicate again before we timeout */
            fprintf(stderr, "%s set PUSH flag after sending %u bytes of payload, hence we're done for now!\n",
            curr_session->ip_src_s, curr_session->bytes_count);
        }
    }

    curr_session->packet_count++;
    return 0;
}

void session_stats(tcpflow_t *list)
{
    tcpflow_t *iter = list;
    puts("SRC IP/PORT              DST IP/PORT             BYTES   PACKETS");
    while (iter)
    {
        printf("%s/%u \t %s/%u \t %u \t %u\n", iter->ip_src_s, iter->src_port,
        iter->ip_dst_s, iter->dst_port, iter->bytes_count, iter->packet_count);
        iter = iter->next;
    }
}

void cleanup_session(tcpflow_t *list)
{
    tcpflow_t *iter = list;
    tcpflow_t *tofree = NULL;
    while (iter)
    {
        close(iter->fd);
        tofree = iter;
        iter = iter->next;
        free(tofree);
    }
}

void handle_packet(u_char *user, const struct pcap_pkthdr *h, const u_char *bytes)
{
    struct ip *iphdr = (struct ip *) (bytes + 14); /* skip ethernet header */
    if (iphdr->ip_p == IPPROTO_TCP)
    {
        struct tcphdr *tcphdr = (struct tcphdr *) (iphdr + 1);

        /* check if session exists for this packet */
        tcpflow_t *session_ptr = NULL;
        if (check_session(head, iphdr->ip_src, iphdr->ip_dst, ntohs(tcphdr->th_sport),
            ntohs(tcphdr->th_dport), &session_ptr) == 0)
            cont_session(session_ptr, iphdr, tcphdr);
        else
            add_session(&head, iphdr, tcphdr);
    }
    else
    {
        fprintf(stderr, "skipping non-tcp packet, protocol was %s\n",
        iphdr->ip_p == IPPROTO_UDP ? "UDP" : "unknown");
    }
}

int main(int argc, char **argv)
{
    if (argc != 3)
    {
        printf("usage: %s <tracefile> <logs>\n", argv[0]);
        exit(1);
    }
    basedir = argv[2];

    struct stat st;
    if (stat(argv[2], &st) == -1)
    {
        fprintf(stderr, "creating directory \"%s\"\n", argv[2]);
        mkdir(argv[2], 0755);
    }

    char errbuf[PCAP_ERRBUF_SIZE];
    pcap_t *handle = pcap_open_offline(argv[1], errbuf);
    while (pcap_loop(handle, -1, &handle_packet, (unsigned char *) "me") > 0)
    {
    }

    session_stats(head);
    cleanup_session(head);

    return 0;
}