#include "../include/fireflier.h"
#include <iostream>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <sys/stat.h>
#include <fcntl.h>
#include "rules.h"

char *CHAINS[]={"", "INPUT", "FORWARD", "OUTPUT"};
static char *IPTABLES[3]={"/usr/sbin/iptables", "/sbin/iptables", "/usr/local/sbin/iptables"};
int iptables_path=-1; // index in IPTABLES

using namespace std;

int initrules()
{
    struct stat filedes;
    iptables_path=0;
    while (iptables_path<3)
    {
	if (stat(IPTABLES[iptables_path], &filedes)!=-1)
	{
            return 1; // success
	}
        iptables_path++;
	}
    iptables_path=-1;
    return 0; // failure
}

// get position of queue target in this chain. (sothat we can insert new rules immediately before)
// if there are more than one queue targets, we use the first one
int getPos(int chainid)
{
    int counter=-1;
    char *rules=getChains(); // get all iptables rules
    char line[150];
    char *pos, *nlpos;

    // valid chain ?
    if ((chainid>=4) || (chainid<=0))
    {
        cout << "RULES LINE " << __LINE__ << ": INVALID CHAIN ID" << endl;
        return -1;
    }

    if (!rules)
        return -1;

    pos=rules;
    // now I parse the rules. as soon as beginning of desired chain is reached, I start a counter
    // count until queue rule reached --> return counter
    do {
        nlpos=strchr(pos, '\n'); // divide into lines
        if (nlpos)
        {
            if (nlpos-pos>150) // just for sure
                return -1;
            strncpy(line, pos, nlpos-pos); 
            line[nlpos-pos]=0;

            if (strstr(line, CHAINS[chainid])!=NULL) // beginning of desired chain?
            {
                counter=0;
            }
            else if ((strstr(line, "QUEUE")!=NULL) && (counter!=-1)) // queue target reached
            {
                nlpos=0; // to exit the loop
            }
            else if (counter>=0) // counting activated?
            {
                counter++; 
            }
            pos=nlpos+1;
        }
    } while (nlpos);
    delete[] rules;
    return counter;
}

// input a rule by various parameters. if any parameter is 0, it is ignored
int inputRule(unsigned int chainid, int protocol, unsigned long int srcip, unsigned long int dstip, int srcport, int dstport, char *interface_in, char *interface_out, int conntrack, int action)
{
    char cmd[300];
    char revcmd[300];
    char temp[30];
    int pos;

    // should never happen
    if (iptables_path==-1)
        return 0;

    // valid rule ?
    if ((action!=ff::ACTION_QUEUE) && !protocol && !srcip && !dstip && !srcport && !dstport && !interface_in &&
        !interface_out)
        return 0;

    // valid chain ?
    if ((chainid>=4) || (chainid<=0))
    {
        cout << "RULES LINE " << __LINE__ << ": INVALID CHAIN ID" << endl;
        return 0;
    }
    
    // inserting of QUEUE rule?
    if (action==ff::ACTION_QUEUE) // create queue rule as last rule in firewall
    {
        strcpy(cmd, IPTABLES[iptables_path]);
        strcat(cmd, " -A ");
        strcpy(revcmd, IPTABLES[iptables_path]);
        strcat(revcmd, " -A ");
    }
    else // others create at position queue-1 (calculation of position follows
    {
        strcpy(cmd, IPTABLES[iptables_path]);
        strcat(cmd, " -I ");
        strcpy(revcmd, IPTABLES[iptables_path]);
        strcat(revcmd, " -I ");
    }
    strcat(cmd, CHAINS[chainid]); // chain identifier
    strcat(revcmd, CHAINS[4-chainid]);
    
    if (action!=ff::ACTION_QUEUE) // for ff::ACTION_QUEUE not needed (append)
    {
        pos=getPos(chainid); // position queue-1
        if (pos==-1) // should not happen
        {
            cout << "Position not found!!!" << endl;
            return 0;
        }
        sprintf(temp, " %d ", pos);
        strcat(cmd, temp);
        
        pos=getPos(4-chainid); // position queue-1
        if (pos==-1) // should not happen
        {
            cout << "Position not found!!!" << endl;
            return 0;
        }
        sprintf(temp, " %d ", pos);
        strcat(revcmd, temp);
    }

    if (protocol!=0) // handle protocols (number --> protocol string)
    {
        switch (protocol)
        {
	case ff::PROT_TCP:
            strcat(cmd, "-p tcp");
            strcat(revcmd, "-p tcp");
            break;
	case ff::PROT_UDP:
            strcat(cmd, "-p udp");
            strcat(revcmd, "-p udp");
            break;
	case ff::PROT_ICMP:
            strcat(cmd, "-p icmp");
            strcat(revcmd, "-p icmp");
            break;
	case ff::PROT_GRE:
            strcat(cmd, "-p gre");
            strcat(revcmd, "-p gre");
            break;
        default:
            break;
        }
    }

    if (srcip!=0) // match Source ip of packets
    {
        strcat(cmd, " -s ");
        strcat(revcmd, " -d ");
        sprintf(temp, "%d.%d.%d.%d", ((unsigned char*)&srcip)[0], ((unsigned char*)&srcip)[1], ((unsigned char*)&srcip)[2], ((unsigned char*)&srcip)[3]);
        strcat(cmd, temp);
        strcat(revcmd, temp);
    }
    if (dstip!=0) // match Destination ip of packets
    {
        strcat(cmd, " -d ");
        strcat(revcmd, " -s ");
        sprintf(temp, "%d.%d.%d.%d", ((unsigned char*)&dstip)[0], ((unsigned char*)&dstip)[1], ((unsigned char*)&dstip)[2], ((unsigned char*)&dstip)[3]);
        strcat(cmd, temp);
        strcat(revcmd, temp);
    }

    if (srcport!=0) // match source port
    {
        sprintf(temp, " --source-port %d", srcport);
        strcat(cmd, temp);
//        sprintf(temp, " --destination-port %d", srcport);
//        strcat(revcmd, temp);
    }
    if (dstport!=0) // match destination port
    {
        sprintf(temp, " --destination-port %d", dstport);
        strcat(cmd, temp);
//        sprintf(temp, " --source-port %d", dstport);
//        strcat(revcmd, temp);
    }

    if (interface_in!=0) // match incoming interface names
    {
        sprintf(temp, " --in-interface %s", interface_in);
        strcat(cmd, temp);
        sprintf(temp, " --out-interface %s", interface_in);
        strcat(revcmd, temp);
    }
    if (interface_out!=0) // match outgoing interface names
    {
        sprintf(temp, " --out-interface %s", interface_out);
        strcat(cmd, temp);
        sprintf(temp, " --in-interface %s", interface_out);
        strcat(revcmd, temp);
    }

    // stateful firewalling (allow only those packets in the other direction,
    // which belong to existing connection). see connection tracking
    strcat (revcmd, " -m state --state RELATED,ESTABLISHED");

    if (action==ff::ACTION_DROP)
    {
        strcat(cmd, " -j DROP");
        strcat(revcmd, " -j DROP");
    }
    else if (action==ff::ACTION_ACCEPT)
    {
        strcat(cmd, " -j ACCEPT");
        strcat(revcmd, " -j ACCEPT");
    }
    else if (action==ff::ACTION_QUEUE)
    {
        strcat(cmd, " -j QUEUE");
        strcat(revcmd, " -j QUEUE");
    }

//    printf("%s\n", cmd);
    if (conntrack) // only if connection tracking wanted use inverted rule
    {
		system(revcmd);
    }
    system(cmd); // insert rule
    return 1;
}

// delete rule by definition. needed by timeoutthread
void deleteRule(unsigned int chainid, int protocol, unsigned long int srcip, unsigned long int dstip, int srcport, int dstport, char *interface_in, char *interface_out, int conntrack, int action)
{
    char cmd[300];
    char revcmd[300];
    char temp[30];
//    char *rules;

    // valid rule?
    if ((action!=ff::ACTION_QUEUE) && !protocol && !srcip && !dstip && !srcport && !dstport && !interface_in &&
        !interface_out)
        return;
   
    // valid chain?
    if ((chainid>=4) || (chainid<=0))
    {
        cout << "RULES LINE " << __LINE__ << ": INVALID CHAIN ID" << endl;
        return;
    }


    strcpy(cmd, IPTABLES[iptables_path]);
    strcat(cmd, " -D ");
    strcpy(revcmd, IPTABLES[iptables_path]);
    strcat(revcmd, " -D ");
    strcat(cmd, CHAINS[chainid]);
    strcat(revcmd, CHAINS[4-chainid]);
    strcat(cmd, " ");
    strcat(revcmd, " ");

    if (protocol!=0) // protocols (number --> name)
    {
        switch (protocol)
        {
	case ff::PROT_TCP:
            strcat(cmd, "-p tcp");
            strcat(revcmd, "-p tcp");
            break;
        case ff::PROT_UDP:
            strcat(cmd, "-p udp");
            strcat(revcmd, "-p udp");
            break;
        case ff::PROT_ICMP:
            strcat(cmd, "-p icmp");
            strcat(revcmd, "-p icmp");
            break;
        case ff::PROT_GRE:
            strcat(cmd, "-p gre");
            strcat(revcmd, "-p gre");
            break;
        default:
            break;
        }
    }

    if (srcip!=0) // source of packets
    {
        strcat(cmd, " -s ");
        strcat(revcmd, " -d ");
        sprintf(temp, "%d.%d.%d.%d", ((unsigned char*)&srcip)[0], ((unsigned char*)&srcip)[1], ((unsigned char*)&srcip)[2], ((unsigned char*)&srcip)[3]);
        strcat(cmd, temp);
        strcat(revcmd, temp);
    }
    if (dstip!=0) // destination of packets
    {
        strcat(cmd, " -d ");
        strcat(revcmd, " -s ");
        sprintf(temp, "%d.%d.%d.%d", ((unsigned char*)&dstip)[0], ((unsigned char*)&dstip)[1], ((unsigned char*)&dstip)[2], ((unsigned char*)&dstip)[3]);
        strcat(cmd, temp);
        strcat(revcmd, temp);
    }

    if (srcport!=0) // source port of packets
    {
        sprintf(temp, " --source-port %d", srcport);
        strcat(cmd, temp);
//        sprintf(temp, " --destination-port %d", srcport);
//        strcat(revcmd, temp);
    }
    if (dstport!=0) // destination port of packets
    {
        sprintf(temp, " --destination-port %d", dstport);
        strcat(cmd, temp);
//        sprintf(temp, " --source-port %d", dstport);
//        strcat(revcmd, temp);
    }

    if (interface_in!=0) // incoming interface
    {
        sprintf(temp, " --in-interface %s", interface_in);
        strcat(cmd, temp);
        sprintf(temp, " --out-interface %s", interface_in);
        strcat(revcmd, temp);
    }
    if (interface_out!=0) // ougoing interface
    {
        sprintf(temp, " --out-interface %s", interface_out);
        strcat(cmd, temp);
        sprintf(temp, " --in-interface %s", interface_out);
        strcat(revcmd, temp);
    }

    // connection tracking
    strcat (revcmd, " -m state --state RELATED,ESTABLISHED");

    if (action==ff::ACTION_DROP)
    {
        strcat(cmd, " -j DROP");
        strcat(revcmd, " -j DROP");
    }
    else if (action==ff::ACTION_ACCEPT)
    {
        strcat(cmd, " -j ACCEPT");
        strcat(revcmd, " -j ACCEPT");
    }
    else if (action==ff::ACTION_QUEUE)
    {
        strcat(cmd, " -j QUEUE");
        strcat(revcmd, " -j QUEUE");
    }

    system(cmd);

    // corresponding inverted rule if connection tracking is enabled
    if (conntrack)
    {
        system(revcmd);
    }
}

// delete rule by position in chain
void deleteRule(unsigned int chainid, int rulenum)
{
    char rule[8];
    char cmd[80];

    //should never happen
    if (iptables_path==-1)
        return;

    // valid chain?
    if ((chainid>3) || (chainid<1)) return;

    strcpy(cmd, IPTABLES[iptables_path]);
    strcat(cmd, " -D ");
    strcat(cmd, CHAINS[chainid]);

    sprintf(rule, " %d ", rulenum);
    strcat(cmd, rule);

    system(cmd);
}

// get all firewall chains as returned by iptables -L -n -v --line-numbers
// this output is used for the interface and for finding rules in the chains
char *getChains()
{
    char *buffer=0;
    int filesize;
    struct stat filedes;
    int handle;
    char cmd[200];

    // create file with output of iptables (I use this as it is easier than using libiptc
    // and iptables itself should work more reliable than implementing this functionality myself)
    strcpy(cmd, "rm -f /var/run/fireflier/fireflier.rules && touch /var/run/fireflier/fireflier.rules && chmod 0700 /var/run/fireflier/fireflier.rules && ");
    strcat(cmd, IPTABLES[iptables_path]);
    strcat(cmd, " -v -L -n --line-numbers > /var/run/fireflier/fireflier.rules");
    system(cmd);

    // should usually work
    if (stat("/var/run/fireflier/fireflier.rules", &filedes)==-1)
        return 0;

    filesize=filedes.st_size;

    buffer=new char[filesize+1];
    if (!buffer)
    {
        cout << "Out of memory" << endl;
        return 0;
    }
    buffer[0]=0;
    buffer[filesize]=0;

    // open rules file
    if ((handle=open("/var/run/fireflier/fireflier.rules", O_RDONLY))==-1)
        return 0;

    if (read(handle, buffer, filesize)==-1)
        return 0;

    close(handle);
    // finally delete the rules file to clean up
    unlink("/var/run/fireflier/fireflier.rules");

    return buffer;
}