#include "sockets/sslsocket.h"
#include "args.h"
#include "util/logmsg.h"

#include "sockets/selectloop.h"
#include "sockets/sslstate.h"
#include "sockets/socks4state.h"
#include "sockets/tcpstate.h"

// todo: fix ssl over tor
int g_connecttimeout= 5000; // msec
int g_sessiontimeout= 5000; // msec
int g_pollinterval= 1; // seconds
int g_verbose= 0;
unsigned g_maxsimultaneous= 32;

class addresslist {
    StringList _targets;
    StringList _ports;
    StringList _addrs;

    size_t _ia;
    size_t _ip;
public:
    addresslist()
        : _ia(0), _ip(0)
    {
    }
    void addports(const std::string& ports)
    {
        _ports.push_back(ports);
    }
    void addaddr(const std::string& addr)
    {
        _addrs.push_back(addr);
    }
    void reset()
    {
        _ia= 0;
        _ip= 0;
    }
    void expand()
    {
        StringList expandedports;
        for (StringList::iterator i= _ports.begin() ; i!=_ports.end() ; ++i)
            expandlist(*i, expandedports);

        _ports.swap(expandedports);

        StringList expandedaddrs;
        for (StringList::iterator i= _addrs.begin() ; i!=_addrs.end() ; ++i)
            if (!expandip(*i, expandedaddrs))
                expandedaddrs.push_back(*i);

        _addrs.swap(expandedaddrs);

        printf("%d ports, %d addrs\n", (int)_ports.size(), (int)_addrs.size());
    }
    // transform "first-last" -> [ first, first+1, first+2, ... last ]
    bool expandrange(const std::string& rangespec, StringList& expanded)
    {
        char *p;
        int first= strtol(&rangespec[0], &p, 0);
        if (p==&rangespec[0] || *p!='-')
            return false;
        p++;
        char *q;
        int last= strtol(p, &q, 0);
        if (p==q || *q)
            return false;
        for (int i=first ; i<=last ; i++)
            expanded.push_back(stringformat("%d", i));
        return true;
    }

    bool isnumber(const std::string& str)
    {
        return str.find_first_not_of("0123456789")==str.npos;
    }
    // transform "<elem>,<elem>,<elem>"  ->  [ "<elem>", "<elem>", "<elem>" ]
    void expandlist(const std::string& listspec, StringList& expanded)
    {
        StringList l;
        if (!SplitString(listspec, l, false, ","))
            throw "invalid list";
        for (StringList::iterator i= l.begin() ; i!=l.end() ; ++i)
            if (!expandrange(*i, expanded))
            {
                if (!isnumber(*i))
                    throw "not a number";
                expanded.push_back(*i);
            }
    }
    bool incnum(std::vector<unsigned>& p, const std::vector<unsigned>& m) 
    {
        unsigned i=0;
        while (i<p.size() && p[i]==m[i])
            p[i++]= 0;
        if (i<p.size()) {
            p[i]++;
            return true;
        }
        return false;
    }

    // transform "1.2.3,8.0-255" -> ["1.2.3.0", "1.2.3.1", ..., "1.2.3.255", "1.2.8.0", ... "1.2.8.255" ]
    bool expandip(const std::string& addrspec, StringList& expanded)
    {
        StringList l;
        if (!SplitString(addrspec, l, false, ".")
                || l.size()!=4 ) {
            return false;
        }
        std::vector<StringList> n(4);
        for (unsigned i= 0 ; i<4 ; i++)
            if (!expandrange(l[i], n[i])) {
                if (!isnumber(l[i]))
                    return false;
                n[i].push_back(l[i]);
            }

        std::vector<unsigned> m(4);
        for (unsigned i= 0 ; i<4 ; i++)
            m[i]= n[i].size()-1;
        std::vector<unsigned> p(4);
        do {
            expanded.push_back(n[0][p[0]]+"."+n[1][p[1]]+"."+n[2][p[2]]+"."+n[3][p[3]]+".");
        } while (incnum(p, m));
        return true;
    }
    std::string next()
    {
        if (!havemore())
            throw "eof";
        std::string addr= _addrs[_ia]+":"+_ports[_ip];
        if (_ia<_addrs.size()-1) {
            _ia++;
        }
        else if (_ia==_addrs.size()-1) {
            _ia= 0;
            if (_ip<_ports.size()-1)
                _ip++;
            else {
                //printf("%s last\n", logstamp().c_str());
                _ia= _addrs.size();
                _ip= _addrs.size();
            }
        }
        return addr;
    }
    bool havemore()
    {
        return _ia<_addrs.size() && _ip<_ports.size();
    }
};

enum { MODE_TCP, MODE_TOR };
void addtarget(int mode, selectloop& l, sslcontext_ptr ctx, const std::string& target)
{
    socket_ptr s;
    socket_ptr ssl;
    socket_ptr socks;
    if (ctx) {
        ssl= socket_ptr(new sslstate(ctx));
        ssl->connecttimeout(g_connecttimeout);
        ssl->sessiontimeout(g_sessiontimeout);
        l.add(ssl);
    }

    if (mode==MODE_TCP) {
        s= socket_ptr(new tcpstate(tcpaddress(target)));
        if (ctx)
            s->next(ssl);
    }
    else if (mode==MODE_TOR) {
        s= socket_ptr(new tcpstate(tcpaddress("127.0.0.1:9050")));


        socks= socket_ptr(new socks4state(target));
        socks->connecttimeout(g_connecttimeout*3);
        socks->sessiontimeout(g_sessiontimeout);
        s->next(socks);

        if (ctx)
            socks->next(ssl);

        l.add(socks);
    }
    s->verbose(g_verbose);
    s->connecttimeout(g_connecttimeout);
    s->sessiontimeout(g_sessiontimeout);
    l.add(s);

    (ctx ? ssl : socks ? socks : s)->write(stringformat("GET / HTTP/1.0\r\nHost: %s\r\n\r\n", target.c_str()));

    s->start(socket_ptr());

    l.bump();
}
void usage()
{
    printf("Usage: torscan [options] [-p PORT ...] [ADDR(s)]\n");
    printf("   -t       scan via TOR\n");
    printf("   -s       try ssl\n");
    printf("   -T msec  connect timeout\n");
    printf("   -P sec   timeout check interval\n");
    printf("   -C MAX   max simultaneous connections\n");
}
int main(int argc, char**argv)
{
    int mode= MODE_TCP;
    bool usessl= false;
    addresslist targets;
    for (int i=1 ; i<argc ; i++)
    {
        if (argv[i][0]=='-') switch(argv[i][1])
        {
            case 'p': targets.addports(getstrarg(argv, i, argc)); break;
            case 't': mode=MODE_TOR; break;
            case 's': usessl= true; break;
            case 'T': g_connecttimeout= getintarg(argv, i, argc); break;
            case 'S': g_sessiontimeout= getintarg(argv, i, argc); break;
            case 'P': g_pollinterval= getintarg(argv, i, argc); break;
            case 'C': g_maxsimultaneous= getintarg(argv, i, argc); break;
            case 'v': g_verbose++; break;
            default:
                      usage();
                      return 1;
        }
        else {
            targets.addaddr(argv[i]);
        }
    }

    try {
    selectloop l;
    l.verbose(g_verbose);
    l.pollinterval(g_pollinterval);

    sslcontext_ptr ctx;
    if (usessl)
        ctx= sslcontext_ptr(new sslcontext(false));

    targets.expand();

    // do scan
    while (l.count() || targets.havemore()) {

        while (targets.havemore() && l.count() < g_maxsimultaneous)
            addtarget(mode, l, ctx, targets.next());

        sleep(1);
    }

    }
    catch(...)
    {
        printf("%s EXCEPTION in main\n", logstamp().c_str());
    }

    return 0;
}