#include <stdio.h>
#include <stdlib.h>
#include <cstring>
#include <unistd.h>
#include <sys/un.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <netinet/ip.h>
#include <netinet/in.h>
#include <netinet/udp.h>
#include <netinet/ip_icmp.h>
#include <netdb.h>
#include <iostream>
#include <string.h>
#include <vector>
#include <regex>
#include <sstream>

// ---------- HELPER STRUCTS ----------
struct port_message {
    int port;
    std::string message;
};

struct pseudo_header {
    u_int32_t source_address;
    u_int32_t dest_address;
    u_int8_t placeholder;
    u_int8_t protocol;
    u_int16_t udp_length;
};
// ---------- GLOBAL VARIABLES ----------
#define BUFFERSIZE 1024
std::vector<port_message> ports;
int socketDGRM;
int socketRAW;
int socketICMP;
int srcPort = 5600;
struct sockaddr_in client;
struct sockaddr_in serverAddress;
char srcIP[32] = "192.168.1.5";

// ---------- SMALL HELPER FUNCTIONS ----------

// get local ip, source: https://bit.ly/2kiprB4
void getLocalIp(char * buffer, int socketDGRM) {
	socketDGRM = socket(AF_INET, SOCK_DGRAM, 0);

	const char* kGoogleDnsIp = "8.8.8.8";
	int dns_port = 53;

	struct sockaddr_in serv;

	memset(&serv, 0, sizeof(serv));
	serv.sin_family = AF_INET;
	serv.sin_addr.s_addr = inet_addr(kGoogleDnsIp);
	serv.sin_port = htons(dns_port);

	int err = connect(socketDGRM,(const struct sockaddr*) &serv , sizeof(serv));

	struct sockaddr_in name;
	socklen_t namelen = sizeof(name);
	err = getsockname(socketDGRM, (struct sockaddr*) &name, &namelen);

	const char *p = inet_ntop(AF_INET, &name.sin_addr, buffer, 100);
	close(socketDGRM);
}

// Creates a new DGRAM socket
void createSocketDGRM() {
    if((socketDGRM = socket(AF_INET , SOCK_DGRAM , IPPROTO_UDP)) == -1) {
        perror("creating udp socket failed");
        exit(0);
    }
}
// Creates a new RAW socket
void createSocketRAW() {
    if((socketRAW = socket(AF_INET , SOCK_RAW , IPPROTO_RAW)) == -1) {
        perror("creating raw socket failed");
        exit(0);
    }
    int on = 1;
    if((setsockopt(socketRAW, IPPROTO_IP, IP_HDRINCL, &on, sizeof(on))) == -1) {
        perror("setsockopt on raw socket failed");
        exit(0);
    }
}
// Creates a new ICMP socket
void createSocketICMP() {
    if((socketICMP = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP)) == -1) {
        perror("creating raw socket failed");
        exit(0);
    }
}
// Binds given socket
void bindSocket(int socket, struct sockaddr_in s) {
    if(bind(socket, (struct sockaddr*)&s, sizeof(s)) == -1) {
        perror("binding socket failed");
        exit(0);
    }
}
// Assigns addresses for both sending and receiving
void assignAddress(char *argv[]) {
    serverAddress.sin_family = AF_INET;
    serverAddress.sin_addr.s_addr = inet_addr(argv[1]);

    client.sin_family = AF_INET;
    client.sin_addr.s_addr = INADDR_ANY;
    client.sin_port = htons(srcPort);
}
// Extracts port from given message
int extractPort(std::string message) {
    int i;
    for (i = 0; i < message.length(); i++ ){
        if ( isdigit(message[i]) ) {
            break;
        }
    }
    message = message.substr(i, message.length() - i);
    return atoi(message.c_str());
}
// Extracts secret phrase from message
std::string extractSecretPhrase(std::string message) {
    char buffer[256];
    char *subString;

    strcpy(buffer, message.c_str());

    subString = strtok(buffer,"\"");
    subString = strtok(NULL,"\"");

    return subString;
}
// Calculates the checksum for a custom packet
unsigned short csum(unsigned short *ptr,int nbytes) {
    unsigned long sum = 0;
    unsigned short oddbyte;
    unsigned short answer;

    while(nbytes>1) {
        sum+=(*ptr++);
        nbytes-=2;
    }
    if(nbytes==1) {
        oddbyte=0;
        *((u_char*)&oddbyte)=*(u_char*)ptr;
        sum+=(oddbyte);
    }
    sum = (sum>>16)+(sum & 0xffff);
    sum = sum + (sum>>16);
    answer=(short)~sum;
    return (answer);
}
// Calculates data for packet from two checksums
unsigned short desiredData(unsigned short currentCsum, int dSum) {
    unsigned short desiredCsum = htons(dSum);
    unsigned short desiredTsum = ~desiredCsum;
    unsigned short currentTsum = ~currentCsum;
    unsigned short diff = desiredTsum - current Tsum;
    return htons(diff);
}
// ---------- SEND AND RECEIVE FUNCTIONS ----------
// sends package to given port
void sendToPort(int socket, struct sockaddr_in serverAddress, const char *package, size_t packageSize) {
    // sending message to server
    if (sendto(socket, package, packageSize, 0, (const struct sockaddr *)&serverAddress, sizeof(serverAddress)) == -1) {
        perror("sendto failed");
        exit(0);
    }
}
// Listens to given port for UDP packets for a certain amount of time
port_message receiveFromPort(int socket, struct sockaddr_in serverAddress, int port) {
    // setting timeout limit
    struct timeval time;
    time.tv_sec = 0;
    time.tv_usec = 200000;
    fd_set fdSet;
    FD_ZERO(&fdSet);
    FD_SET(socket, &fdSet);
    socklen_t recvAddrSize;
    char buffer[BUFFERSIZE];
    if(select(socket + 1, &fdSet, NULL, NULL, &time) > 0) {
        memset(buffer, 0, BUFFERSIZE);
        // waiting to receive message from server
        if(recvfrom(socket, buffer, sizeof(buffer), 0, (struct sockaddr *)&serverAddress, &recvAddrSize) == -1) {
            perror("recvfrom failed");
            exit(0);
        }
        port_message m;
        m.port = port;
        m.message = buffer;
        return m;
    }
    port_message m;
    m.port = port;
    m.message = "no answer from port";
    return m;
}
// Listens to given port for ICMP packets for a certain amount of time
port_message receiveICPMFromPort(int socket, struct sockaddr_in serverAddress, int port) {
    // setting timeout limit
    struct timeval time;
    time.tv_sec = 0;
    time.tv_usec = 200000;
    fd_set fdSetICMP;
    FD_ZERO(&fdSetICMP);
    FD_SET(socket, &fdSetICMP);
    socklen_t recvAddrSize;
    char buffer[BUFFERSIZE];
    if(select(socket + 1, &fdSetICMP, NULL, NULL, &time) > 0) {
        memset(buffer, 0, BUFFERSIZE);
        // waiting to receive message from server
        if(recvfrom(socket, buffer, sizeof(buffer), 0, (struct sockaddr *)&serverAddress, &recvAddrSize) == -1) {
            perror("recvfrom failed");
            exit(0);
        }
        struct ip *ipheader = (struct ip*)buffer;
        int iplen = ipheader->ip_hl << 2;
        struct icmp *icmp = (struct icmp*)(buffer + iplen);
        if((icmp->icmp_type == ICMP_UNREACH) && (icmp->icmp_code == ICMP_UNREACH_PORT)) {
            port_message m;
            m.port = port;
            m.message = "unreachable port";
            return m;
        }
    }
    port_message m;
    m.port = port;
    m.message = "no answer from port";
    return m;
}
// ---------- UDP PORT SCANNER ----------
// Scans given port range and saves their messages
void scanPorts(int lowPort, int highPort) {
    // assigning sockets
    createSocketDGRM();
    createSocketICMP();
    const char *package = "knock";
    port_message m1, m2;
    int count = 0;
    // checking all the ports in port range
    for(int i = lowPort; i <= highPort; i++) {
        serverAddress.sin_port = htons(i);
        sendToPort(socketDGRM, serverAddress, package, strlen(package));
        m2 = receiveICPMFromPort(socketICMP, serverAddress, i);
        if(m2.message == "no answer from port") {
            m1 = receiveFromPort(socketDGRM, serverAddress, i);
            if(m1.message != "no answer from port" || count > 3) {
                if(m1.message.find("oracle")) {
                    ports.insert(ports.begin(), m1);
                } else {
                    ports.push_back(m1);
                }
                count = 0;
            }
            else {
                count++;
                i--;
            }
        }
        else {
            count = 0;
        }
    }
    std::cout << "---- all ports in range ----" << std::endl;
    for(port_message p : ports) {
        std::cout << "port: " << p.port << " | " << p.message << std::endl;
    }
    close(socketDGRM);
    close(socketICMP);
}
// ---------- FUNCTIONS FOR MAKING PACKETS ----------
// Creates a custom udp packet with given variables
char* createPacket(struct sockaddr_in serverAddress, int dstPort, char *myData, u_short ipOff) {
    char packet[4096], *data, *pseudogram;
    memset (packet, 0, 4096);
    // IP header of packet
    struct ip *iph = (struct ip *) packet;
    // UDP header of packet
    struct udphdr *udph = (struct udphdr *) (packet + sizeof (struct ip));
    // Data of packet
    data = packet + sizeof(struct ip) + sizeof(struct udphdr);
    strcpy(data , myData);
    // Assigning IP header
    iph->ip_hl = 5;
    iph->ip_v = 4;
    iph->ip_tos = 0;
    iph->ip_len = sizeof (struct ip) + sizeof (struct udphdr) + strlen(data);
    iph->ip_id = 0;
    iph->ip_off = ipOff;
    iph->ip_ttl = 255;
    iph->ip_p = IPPROTO_UDP;
    iph->ip_sum = 0;
    iph->ip_src.s_addr = inet_addr (srcIP);
    iph->ip_dst.s_addr = serverAddress.sin_addr.s_addr;
    // Assigning UDP header
    udph->uh_sport = htons(srcPort);
    udph->uh_dport = htons(dstPort);
    udph->uh_ulen = htons(8 + strlen(data));
    udph->uh_sum = 0;
    // Creating pseudo header for calculating checksum
    struct pseudo_header psh;
    psh.source_address = inet_addr(srcIP);
    psh.dest_address = serverAddress.sin_addr.s_addr;
    psh.placeholder = 0;
    psh.protocol = IPPROTO_UDP;
    psh.udp_length = htons(sizeof(struct udphdr) + strlen(data));
    int psize = sizeof(struct pseudo_header) + sizeof(struct udphdr) + strlen(data);
    pseudogram = (char *)malloc(psize);
    memcpy(pseudogram , (char*) &psh , sizeof (struct pseudo_header));
    memcpy(pseudogram + sizeof(struct pseudo_header) , udph , sizeof(struct udphdr) + strlen(data));
    // Calculate udp checksum
    udph->uh_sum = csum( (unsigned short*) pseudogram , psize);
    return packet;
}
// Creates packet with empty data and calculates data from its checksum and a given checksum
unsigned short makeDataForPacket(struct sockaddr_in serverAddress, int dstPort, int dSum) {
    char packet[4096], *data, *pseudogram;
    memset (packet, 0, 4096);
    // UDP header of packet
    struct udphdr *udph = (struct udphdr *) (packet + sizeof (struct ip));
    // Data of packet
    data = packet + sizeof(struct ip) + sizeof(struct udphdr);
    strcpy(data , "");
    // Assigning udp header
    udph->uh_sport = htons(srcPort);
    udph->uh_dport = htons (dstPort);
    udph->uh_ulen = htons(12 + strlen(data));
    udph->uh_sum = 0;
    // Creating pseudo header for calculating checksum
    struct pseudo_header psh;
    psh.source_address = inet_addr(srcIP);
    psh.dest_address = serverAddress.sin_addr.s_addr;
    psh.placeholder = 0;
    psh.protocol = IPPROTO_UDP;
    psh.udp_length = htons(sizeof(struct udphdr) + strlen(data));
    int psize = sizeof(struct pseudo_header) + sizeof(struct udphdr) + strlen(data);
    pseudogram = (char *)malloc(psize);
    memcpy(pseudogram , (char*) &psh , sizeof (struct pseudo_header));
    memcpy(pseudogram + sizeof(struct pseudo_header) , udph , sizeof(struct udphdr) + strlen(data));
    // Calculate udp checksum
    unsigned short checksum = csum( (unsigned short*) pseudogram , psize);
    return desiredData(checksum, dSum);
}
//---------- FUNCTIONS FOR SENDING TO SPECIFIC PORTS ----------
// Gets port from evil port
int evilPort(int dstPort) {
    serverAddress.sin_port = htons(dstPort);
    char *data = "So you have chosen... death.";
    char *packet = createPacket(serverAddress, dstPort, data, 0x8000);
    port_message m;
    createSocketRAW();
    sendToPort(socketRAW, serverAddress, packet, (sizeof(struct ip) + sizeof(struct udphdr) + strlen(data)));
    createSocketDGRM();
    bindSocket(socketDGRM, client);
     m = receiveFromPort(socketDGRM, client, dstPort);
    std::cout << "---- EVIL PORT ----" << std::endl;
    std::cout << "port: " << m.port << " | " << m.message << std::endl;
    close(socketRAW);
    close(socketDGRM);
    return extractPort(m.message);
}
// Gets secret message from udp checksum port
std::string udpChecksumPort(int dstPort, int desiredChecksum) {
    serverAddress.sin_port = htons(dstPort);
    unsigned short sData = makeDataForPacket(serverAddress, dstPort, desiredChecksum);
    char nData[3];
    nData[0] = sData >> 8;
    nData[1] = sData & 0xfffff;
    nData[2] = '\0';
    char *data = nData;
    char *packet = createPacket(serverAddress, dstPort, data, 0);
    port_message m;
    createSocketRAW();
    sendToPort(socketRAW, serverAddress, packet, (sizeof(struct ip) + sizeof(struct udphdr) + strlen(data)));
    createSocketDGRM();
    bindSocket(socketDGRM, client);
     m = receiveFromPort(socketDGRM, client, dstPort);
    std::cout << "---- UDP CHECKSUM PORT ----" << std::endl;
    std::cout << "port: " << m.port << " | " << m.message << std::endl;
    close(socketRAW);
    close(socketDGRM);
    std::cout << "LALALALLALALALAL" << std::endl;
    return extractSecretPhrase(m.message);
}
// Gets secret ports from the oracle
void oraclePort(int dstPort, int port1, int port2, std::string secretMessage) {
    serverAddress.sin_port = htons(dstPort);
    std::string oracleMessage = std::to_string(port1) + "," + std::to_string(port2);
    const char *package = oracleMessage.c_str();
    std::cout << package << std::endl;
    port_message m;
    createSocketDGRM();
    sendToPort(socketDGRM, serverAddress, package, strlen(package));
     m = receiveFromPort(socketDGRM, serverAddress, dstPort);
    std::cout << "---- ORACLE PORT ----" << std::endl;
    std::cout << "port: " << m.port << " | " << m.message << std::endl;
    /*extracting ports from oracle message*/
    std::vector<int> oraclePorts;
    std::stringstream ss(m.message);
    for (int i; ss >> i;) {
        oraclePorts.push_back(i);
        if (ss.peek() == ',')
            ss.ignore();
    }
    const char *data = "knock";
    for(int &port : oraclePorts) {
        serverAddress.sin_port = htons(port);
        if(&port == &oraclePorts.back()) {
            const char *message = secretMessage.c_str();
            dstPort = port;
            sendToPort(socketDGRM, serverAddress, message, strlen(message));
            std::cout << dstPort << " | " << message << std::endl;
        }
        else {
            sendToPort(socketDGRM, serverAddress, data, strlen(data));
        }
    }
    m = receiveFromPort(socketDGRM, serverAddress, dstPort);
    std::cout << "---- SECRET PORT ----" << std::endl;
    std::cout << "port: " << m.port << " | " << m.message << std::endl;
    close(socketDGRM);
}
// Finds ports in port vector and sends the correct message to each one
void sendToPorts() {
    int thisIsPort, evilP;
    std::string secretMessage;
    for(port_message p: ports) {
        if (p.message.find("evil") != std::string::npos) {
            evilP = evilPort(p.port);
        }
        if(p.message.find("checksum") != std::string::npos) {
            int dSum = extractPort(p.message);
            secretMessage = udpChecksumPort(p.port, dSum);
        }
        if(p.message.find("This is the port") != std::string::npos) {
            thisIsPort = extractPort(p.message);
        }
        if(p.message.find("oracle") != std::string::npos) {
            std::cout << "oracle!" << std::endl;
            oraclePort(p.port, evilP, thisIsPort, secretMessage);
        }
    }
}




int main(int argc, char* argv[]) {
    // what port we want to scan from
    int lowPort = std::stoi (argv[2]);
    // what port we want to scan to
    int highPort = std::stoi (argv[3]);
    //initializing localIp
    //getLocalIp(srcIP, socketDGRM);
    assignAddress(argv);
    // ---- PART 1A ----
    scanPorts(lowPort, highPort);
    std::cout <<" printing ip" << std::endl;
    // for(int i = 0; i < strlen(srcIP); i++) {
    //     std::cout << srcIP[i];
    // }
    // ---- PART 1B AND 1C ----
    sendToPorts();
    return 0;
}