#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/socket.h>

#define MAXN_PORT   65536
#define BUFFER_SIZE 65536

struct save_seq;

struct packet
{
    char *data;
    size_t len, clen;
    unsigned int seq;
};

struct writer
{
    int fd;
    size_t len, writed;
    struct packet *q[10000];
    unsigned int seq;
    int qc, maxqc;
    int finish;
};

struct writer *wr[MAXN_PORT];

struct packet *new_packet(char *buffer, size_t size, size_t csize, unsigned int seq);
struct writer *new_writer(int fd, size_t len, unsigned int seq);
void free_packet(struct packet *pack);
void save_packet(struct writer *w, char *buffer, size_t size, size_t csize, unsigned int seq);
void write_packet(struct writer *w, struct packet *p);
void process_packet(char *buffer, size_t size);
void process_tcp_packet(char *buffer, size_t size);
void writer_update(struct writer *w);
void free_writer(struct writer *w);

int main()
{
    int sock;
    char *buffer;
    ssize_t n;
    struct sockaddr saddr;
    socklen_t slen;
    uid_t uid;
    
    uid = getuid();
    if (setuid(0)) {
        puts("Must run as root");
        return 1;
    }
    
    buffer = (char *) calloc(BUFFER_SIZE, sizeof(char));
    
    sock = socket(AF_INET, SOCK_RAW, IPPROTO_TCP);
    /*
    int val = BUFFER_SIZE;
    setsockopt(sock, SOL_SOCKET, SO_RCVBUF, &val, sizeof(val));
    * */
    if (sock < 1) {
        puts("Socket error");
        return 1;
    }
    seteuid(uid);
    
    while (1) {
        n = recvfrom(sock, buffer, BUFFER_SIZE, 0, &saddr, &slen);
        if (n < 0) {
            puts("Recvfrom error");
            return 1;
        }
        process_packet(buffer, n);
    }
    
    close(sock);
    free(buffer);
    return 0;
}

void process_packet(char *buffer, size_t size)
{
    struct iphdr *iph = (struct iphdr*) buffer;
    if (iph->protocol == SOL_TCP)
        process_tcp_packet(buffer, size);
}

void process_tcp_packet(char *buffer, size_t size)
{
    static int cnt = 0;
    int fd;
    struct iphdr *iph;
    struct tcphdr *tcph;
    char *data, *str_data, *start;
    size_t iplen, tcplen, datalen, mpeglen;
    char tmp_path[80];
    
    iph = (struct iphdr *) buffer;
    iplen = iph->ihl * 4;
    
    tcph = (struct tcphdr *) (buffer + iplen);
    tcplen = tcph->doff * 4;
    
    data = buffer + iplen + tcplen;
    datalen = size - iplen - tcplen;
    
    str_data = strndup(data, datalen);
    if (strstr(str_data, "audio/mpeg")) {
        if (wr[tcph->dest]) {
            puts("HMMM, remove previous writer");
            free_writer(wr[tcph->dest]);
        }
        puts("OK, start writing");
        sprintf(tmp_path, "audio%d.mp3", cnt++);
        fd = open(tmp_path, O_WRONLY | O_CREAT | O_TRUNC, 0664);
        start = strstr(data, "Content-Length:");
        wr[tcph->dest] = new_writer(fd, 
                                    atoi(start + 15), 
                                    ntohl(tcph->seq));
        start = strstr(data, "\r\n\r\n") + 4;
        mpeglen = buffer + size - start;
    } else {
        start = data;
        mpeglen = datalen;
    }
    free(str_data);
    
    if (!wr[tcph->dest])
        return;
        
    if (tcph->seq == wr[tcph->dest]->seq)
        printf("%u", tcph->seq);
    
    save_packet(wr[tcph->dest], start, mpeglen, datalen, ntohl(tcph->seq));
    writer_update(wr[tcph->dest]);
    
    printf("\r %.2f%% ", 100.0 * wr[tcph->dest]->writed / 
                                 wr[tcph->dest]->len);
    printf(" [%d unsorted packages (max %d)]", 
            wr[tcph->dest]->qc,
            wr[tcph->dest]->maxqc);
    fflush(stdout);
    
    
    if (wr[tcph->dest]->finish) {
        free_writer(wr[tcph->dest]);
        wr[tcph->dest] = NULL;
        puts("\nOK, stop writing");
    }
}

struct packet *new_packet(char *buffer, size_t size, size_t csize, unsigned int seq)
{
    struct packet *p = (struct packet *) malloc(sizeof(struct packet));
    p->data = (char *) malloc(size);
    p->len = size;
    p->clen = csize;
    p->seq = seq;
    memcpy(p->data, buffer, size);
    return p;
}

void free_packet(struct packet *pack)
{
    free(pack->data);
    free(pack);
}

void save_packet(struct writer *w, char *buffer, size_t size, 
                 size_t csize, unsigned int seq)
{
    struct packet *p = new_packet(buffer, size, csize, seq);
    w->q[w->qc++] = p;
    if (w->maxqc < w->qc)
        w->maxqc = w->qc;
}

void writer_update(struct writer *w)
{
    int do_write = 1, i;
    while (do_write) {
        do_write = 0;
        for (i = 0; i < w->qc; i++)
            if (w->q[i]->seq == w->seq) {
                do_write = 1;
                write_packet(w, w->q[i]);
                free_packet(w->q[i]);
                w->q[i] = w->q[--w->qc];
                break;
            }
    }
}

void write_packet(struct writer *w, struct packet *p)
{
    write(w->fd, p->data, p->len);
    
    w->writed += p->len;
    w->seq += p->clen;
    
    if (w->writed >= w->len)
        w->finish = 1;
}

void free_writer(struct writer *w)
{
    int i = 0;
    
    /* if some packets still unsorted */
    for (i = 0; i < w->qc; i++)
        free_packet(w->q[i]);
    
    close(w->fd);
    free(w);
}

struct writer *new_writer(int fd, size_t len, unsigned int seq)
{
    struct writer *w = (struct writer *) malloc(sizeof(struct writer));
    w->fd = fd;
    w->len = len;
    w->seq = seq;
    w->writed = 0;
    w->qc = 0;
    w->maxqc = 0;
    w->finish = 0;
    return w;
}