#ifndef __supersocks__socks__
#define __supersocks__socks__

#include <iostream>
#include <stdint.h>
#include <netdb.h>
#include <unistd.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <fcntl.h>
#include <sys/uio.h>
#include <arpa/inet.h>
#include <list>
#include <string>
#include <sys/select.h>


#define MAX_BUFFER 10*1024
/*basic structure */

struct supported_auth{
    bool no_auth = 0;
    bool user_password = 0;
};
enum conditions{
    accepted = 0,
    handshaked = 1,
    connected =3
};

struct remote{
    std::list<struct client>::iterator local;
    int sock;
};

struct client{
    std::list<struct remote>::iterator remote;
    int sock;
    enum conditions cond;
};

/*class define */
class Server{
public:
    int main_sock;
    bool encrption;
    std::list<struct remote> rmt;
    std::list<struct client> clnt;
    char *buffer;

    Server(const char* lcl_port){
        this->buffer = (char*)operator new(MAX_BUFFER);
        addrinfo *addr;
        getaddrinfo(NULL, lcl_port, NULL, &addr);
        addr->ai_family = AF_UNSPEC;
        addr->ai_socktype = SOCK_STREAM;
        addr->ai_protocol = 0;
        set_init();

        main_sock = create_main_socket(addr);
        if(main_sock < 0) std::cerr<<"error on creating main socket"<<std::endl;

        if(bind(main_sock, addr->ai_addr, addr->ai_addrlen)< 0 ) std::cerr<<"error on binding"<<std::endl;
        if(listen(main_sock, backlog) < 0) std::cerr<<"error on listening"<<std::endl;
    }

    void wipe_buffer(ssize_t s){
        memset(buffer, 0 , s);
    }

    bool is_ready(int sock){
        if(FD_ISSET(sock, &set_var)) return true;
        else return false;
    }

    inline void select_block(){
        set_var = set_org;
        int rc = select(maxfd+1, &set_var, NULL, NULL, NULL);
        if(rc < 0 ) std::cerr<<"error on select function"<<std::endl;
    }
    inline void set_add(int sock){
        FD_SET(sock, &set_org);
        if(sock > maxfd) maxfd = sock;
    }
    inline void set_del(int sock){
        FD_CLR(sock, &set_org);
    }
private:
    char *local_port;

    const int backlog = 20;

    fd_set set_org;
    fd_set set_var;

    int maxfd = 0;

    int create_main_socket(addrinfo *addr){
        int sock;
        sock = socket(AF_INET, SOCK_STREAM, 0);
        if(sock <0) {
            return -1;
        }
        /*auto close socket when start a new process*/
        int rc = fcntl(sock, F_SETFD, FD_CLOEXEC);
        if(rc != 0){
            close(sock);
            return -1;
        }
        /*multiuse of addresss*/
        int flags = fcntl(sock, F_GETFL, 0);
        if(flags < 0){
            close(sock);
            return -1;
        }
        if(fcntl(sock, F_SETFL, flags | O_NONBLOCK) < 0){
            close(sock);
            return  -1;
        }
        set_add(main_sock);
        return sock;
    }


    inline void set_init(){
        FD_ZERO(&set_org);
        FD_ZERO(&set_var);
    }

};
#endif /* defined(__supersocks__socks__) */

#include <iostream>
#include "socks.h"
bool decison_making_for_local_packet(std::list<struct client>::iterator it, Server ss);

bool connect(std::list<struct client>::iterator it, Server ss);
bool handshake(std::list<struct client>::iterator it, Server ss);
bool client_forward(std::list<struct client>::iterator it, Server ss);

bool remote_forward(std::list<struct remote>::iterator it, Server ss);

static inline void set_port(sockaddr *addr, uint16_t port){
    ((sockaddr_in*)addr)->sin_port = port;
}

int main(int argc, const char * argv[]) {
    // insert code here...
    if(argc != 2) {
        std::cerr<<"error on argument you take"<<std::endl;
        return -1;
    }
    Server ss (argv[1]);

    while(1){
        ss.select_block();
        /*acception of client*/
        if(ss.is_ready(ss.main_sock)){
            struct client clnt;
            clnt.sock = accept(ss.main_sock, NULL, NULL);
            if(clnt.sock < 0){
                std::cerr<<"error on local client acception"<<std::endl;
                /*dont't close the socket when is below 0*/
            }
            else{
                /*flags control */
                std::cerr<<"accepted"<<std::endl;
                int flags = fcntl(clnt.sock, F_GETFL, 0);
                if(flags < 0){
                    close(clnt.sock);
                    std::cerr<<"error on set local client flags"<<std::endl;
                }
                else if(fcntl(clnt.sock, F_SETFL, flags | O_NONBLOCK) < 0){
                    std::cerr<<"error on set local client flags"<<std::endl;
                    close(clnt.sock);
                }
                clnt.cond = accepted;
                ss.set_add(clnt.sock);
                ss.clnt.push_back(clnt);
            }
        }

        /*when no client online*/
        if(ss.clnt.size() == 0) continue;

        /*client handshake/connect/forward */
        bool dele_client = false;
        decltype(ss.clnt.begin()) del_client;
        auto it_clnt = ss.clnt.begin();
        for(int i = 0; i < ss.clnt.size(); i++){
            if(ss.is_ready(it_clnt->sock)){
                if(!decison_making_for_local_packet(it_clnt, ss)){
                    dele_client = true;
                    del_client = it_clnt;
                }
            }
            if(it_clnt != ss.clnt.end()) it_clnt++;
        }
        if(dele_client){
            if(del_client->cond != connected){
                close(del_client->sock);
                ss.clnt.erase(del_client);
            }
            else{
                std::cerr<<"error 1"<<std::endl;
                close(del_client->sock);
                close(del_client->remote->sock);
                ss.rmt.erase(del_client->remote);
                ss.clnt.erase(del_client);
            }
        }

        /*remote forward*/
        bool dele_remote = false;
        decltype(ss.rmt.begin()) del_remote;
        auto it_rmt = ss.rmt.begin();
        for(int i = 0; i < ss.rmt.size();i++){
            if(ss.is_ready(it_rmt->sock)){
                if(!remote_forward(it_rmt, ss)){
                    dele_remote = true;
                    del_remote = it_rmt;
                }
            }
            if(it_rmt != ss.rmt.end()) it_rmt++;
        }
        if(dele_remote){
            std::cerr<<"error 2"<<std::endl;
            close(del_remote->sock);
            close(del_remote->local->sock);
            ss.clnt.erase(del_remote->local);
            ss.rmt.erase(del_remote);
        }

    }
    return 0;
}

bool decison_making_for_local_packet(std::list<struct client>::iterator it, Server ss){
    if(it->cond == accepted) return handshake(it, ss);
    else if(it->cond == handshaked) return connect(it, ss);
    else if(it->cond == connected) return client_forward(it, ss);
    return false;
}

bool handshake(std::list<struct client>::iterator it, Server ss){
    std::cerr<<"handshake"<<std::endl;
    ss.wipe_buffer(10*sizeof(uint8_t));
    auto rc = recv(it->sock, ss.buffer, 10*sizeof(uint8_t), NULL);
    if(rc < 3) return false;
    struct supported_auth auth;
    uint8_t *version = reinterpret_cast<uint8_t*>( ss.buffer );
    uint8_t *auth_type_length = reinterpret_cast<uint8_t*>(ss.buffer+1);
    uint8_t *auth_beginning = reinterpret_cast<uint8_t*>(ss.buffer+2);
    if(*version != 0x05) return false;
    for(int i = 0; i < static_cast<int>(*auth_type_length); i++){
        switch (*(auth_beginning+i)) {
            case 0x00:
                auth.no_auth = true;
                break;
            case 0x02:
                auth.user_password = true;
            default:
                break;
        }
    }
    /*more auth support*/
    if(auth.no_auth){
        uint8_t response[2];
        response[0] = 0x05;
        response[1] = 0x00;
        rc = send(it->sock, response, 2*sizeof(uint8_t), NULL);
        if(rc != 2) return false;
        else{
            it->cond = handshaked;
            return true;
        }
    }
    else return false;
}

bool client_forward(std::list<struct client>::iterator it, Server ss){
    auto rc = recv(it->sock, ss.buffer, MAX_BUFFER,NULL);
    if(rc <= 0) return false;
    auto sd = send(it->remote->sock, ss.buffer, rc, NULL);
    if(sd <= 0) return false;

    return true;
}

bool connect(std::list<struct client>::iterator it, Server ss){
    ss.wipe_buffer(50*sizeof(uint8_t));
    auto rc = recv(it->sock, ss.buffer, 50*sizeof(uint8_t), NULL);
    addrinfo *addr;
    if(rc < 10) return false;
    if((uint8_t)(*ss.buffer) != 0x05 || (uint8_t)(*(ss.buffer+2)) != 0x00) return false;

    int flags;
    int sock;
    struct remote rmt;
    ssize_t length;
    char *domain_buffer = (char *) operator new(50*sizeof(char));

    switch (*((uint8_t*)(ss.buffer+3))) {
        case 0x01:
            //ipv4
            getaddrinfo(NULL, NULL, NULL, &addr);
            addr->ai_family = AF_INET;
            addr->ai_socktype = SOCK_STREAM;
            addr->ai_protocol = 0;
            memcpy(addr->ai_addr->sa_data, ss.buffer+4, 4*sizeof(uint8_t));
            addr->ai_addrlen = 4;
            addr->ai_addr->sa_len = 4;
            addr->ai_addr->sa_family = AF_INET;
            set_port(addr->ai_addr, *((uint16_t*)(ss.buffer+8)));

            sock = socket(addr->ai_family, addr->ai_socktype, addr->ai_protocol);
            if(sock < 0){
                *(ss.buffer+1) = 0x01;
                send(sock, ss.buffer, rc, NULL);

                return false;
            }
            flags = fcntl(it->sock, F_GETFL, 0);
            if(flags < 0){

                *(ss.buffer+1) = 0x01;
                send(sock, ss.buffer, rc, NULL);
                close(sock);
                return false;
            }
            if(fcntl(sock, F_SETFL, flags | O_NONBLOCK) < 0){

                *(ss.buffer+1) = 0x01;
                send(it->sock, ss.buffer, rc, NULL);
                close(sock);
                return  false;
            }
            freeaddrinfo(addr);

            *(ss.buffer+1) = 0x00;
            send(it->sock, ss.buffer, rc, NULL);
            rmt.sock = sock;
            rmt.local = it;
            ss.rmt.push_back(rmt);
            it->cond =connected;

            return true;

            break;
        case 0x03:
            //domain
            length = *(ss.buffer+4);
            if(length > 50) return false;

            memcpy(domain_buffer, ss.buffer+5, length);

            getaddrinfo(domain_buffer, NULL, NULL, &addr);

            addr->ai_socktype = SOCK_STREAM;
            addr->ai_protocol = 0;

            set_port(addr->ai_addr, *((uint16_t*)(ss.buffer+4+length)));

            sock = socket(addr->ai_family, addr->ai_socktype, addr->ai_protocol);

            if(sock < 0){
                *(ss.buffer+1) = 0x01;
                send(sock, ss.buffer, rc, NULL);

                return false;
            }
            flags = fcntl(it->sock, F_GETFL, 0);
            if(flags < 0){

                *(ss.buffer+1) = 0x01;
                send(sock, ss.buffer, rc, NULL);
                close(sock);
                return false;
            }
            if(fcntl(sock, F_SETFL, flags | O_NONBLOCK) < 0){

                *(ss.buffer+1) = 0x01;
                send(it->sock, ss.buffer, rc, NULL);
                close(sock);
                return  false;
            }

            freeaddrinfo(addr);

            *(ss.buffer+1) = 0x00;
            send(it->sock, ss.buffer, rc, NULL);
            rmt.sock = sock;
            rmt.local = it;
            it->cond =connected;
            ss.rmt.push_back(rmt);

            return true;
            break;
        case 0x04:
            //ipv6
            getaddrinfo(NULL, NULL, NULL, &addr);

            addr->ai_family = AF_INET6;
            addr->ai_socktype = SOCK_STREAM;
            addr->ai_protocol = 0;
            memcpy(addr->ai_addr->sa_data, ss.buffer+4, 16*sizeof(uint8_t));
            addr->ai_addrlen = 4;
            addr->ai_addr->sa_len = 4;
            addr->ai_addr->sa_family = AF_INET6;
            set_port(addr->ai_addr, *((uint16_t*)(ss.buffer+20)));

            sock = socket(addr->ai_family, addr->ai_socktype, addr->ai_protocol);

            if(sock < 0){
                *(ss.buffer+1) = 0x01;
                send(sock, ss.buffer, rc, NULL);

                return false;
            }
            flags = fcntl(it->sock, F_GETFL, 0);
            if(flags < 0){

                *(ss.buffer+1) = 0x01;
                send(sock, ss.buffer, rc, NULL);
                close(sock);
                return false;
            }
            if(fcntl(sock, F_SETFL, flags | O_NONBLOCK) < 0){

                *(ss.buffer+1) = 0x01;
                send(it->sock, ss.buffer, rc, NULL);
                close(sock);
                return  false;
            }

            freeaddrinfo(addr);

            *(ss.buffer+1) = 0x00;
            send(it->sock, ss.buffer, rc, NULL);
            rmt.sock = sock;
            rmt.local = it;
            ss.rmt.push_back(rmt);
            it->cond =connected;

            return true;

            break;
        default:
            return false;
            break;
    }
}

bool remote_forward(std::list<struct remote>::iterator it, Server ss){
    auto rc = recv(it->sock, ss.buffer, MAX_BUFFER,NULL);
    if(rc <= 0) return false;
    auto sd = send(it->local->sock, ss.buffer, rc, NULL);
    if(sd <= 0) return false;
    return true;
}