# Author: Nicolas VERDIER (contact@n1nj4.eu)
# This file is part of pr0cks.
#
# pr0cks is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# pr0cks is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with pr0cks.  If not, see <http://www.gnu.org/licenses/>.

import sys
import StringIO
import time
import struct
import os
import asyncore
import socket
import socks
import argparse
import traceback
import logging
logging.basicConfig(stream=sys.stderr, level=logging.WARNING)
import binascii
from collections import OrderedDict

dnslib_imported=False
dns_cache=OrderedDict()
DNS_CACHE_SIZE=1000

def display(msg):
    msg=msg.strip()
    if msg.startswith("[-]"):
        print "\033[31m[-]\033[0m"+msg[3:]
    elif msg.startswith("[+]"):
        print "\033[32m[+]\033[0m"+msg[3:]
    elif msg.startswith("[i]"):
        print "\033[1;30m[i]\033[0m"+msg[3:]
    else:
        print msg

try:
    from dnslib import DNSRecord, QTYPE
    from dnslib.server import DNSServer,DNSHandler,BaseResolver,DNSLogger
    class ProxyResolver(BaseResolver):
        def __init__(self,address,port):
            self.address = address
            self.port = port

        def resolve(self,request,handler):
            if handler.protocol == 'udp':
                proxy_r = request.send(self.address,self.port)
            else:
                proxy_r = request.send(self.address,self.port,tcp=True)
            reply = DNSRecord.parse(proxy_r)
            return reply

    class PassthroughDNSHandler(DNSHandler):
        def get_reply(self,data):
            global dns_cache
            global args
            host,port = self.server.resolver.address,self.server.resolver.port
            request = DNSRecord.parse(data)


            domain=str(request.q.qname)
            qtype=str(QTYPE.get(request.q.qtype))
            index=domain+"/"+qtype
            if not args.no_cache and index in dns_cache:
                if time.time()<dns_cache[index][0]:
                    if args is not None and args.verbose:
                        try:
                            display("[i] %s served value from cache: %s"%(index, ', '.join([x.rdata for x in dns_cache[index][1]])))
                        except:
                            pass
                    rep=request.reply()
                    rep.add_answer(*dns_cache[index][1])
                    return rep.pack()
            if args is not None and args.verbose:
                display("[i] domain %s requested using TCP server %s"%(domain, args.dns_server))
            data = struct.pack("!H",len(data)) + data
            response = send_tcp(data,host,port)
            response = response[2:]
            reply = DNSRecord.parse(response)
            if args.verbose:
                try:
                    display("[i] %s %s resolve to %s"%(domain, qtype, ', '.join([x.rdata for x in reply.rr])))
                except:
                    pass
            ttl=3600
            try:
                ttl=reply.rr[0].ttl
            except Exception:
                try:
                    ttl=reply.rr.ttl
                except Exception:
                    pass
            dns_cache[index]=(int(time.time())+ttl, reply.rr)
            if len(dns_cache)>DNS_CACHE_SIZE:
                dns_cache.popitem(last=False)
            return response

    def send_tcp(data,host,port):
        """
            Helper function to send/receive DNS TCP request
            (in/out packets will have prepended TCP length header)
        """
        sock = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
        sock.settimeout(5)
        sock.connect((host,port))
        sock.sendall(data)
        response = sock.recv(8192)
        length = struct.unpack("!H",bytes(response[:2]))[0]
        while len(response) - 2 < length:
            response += sock.recv(8192)
        sock.close()
        return response
    dnslib_imported=True
except ImportError:
    display("[-] WARNING: The following dependency is needed to proxify DNS through tcp: pip install dnslib")

args=None
if __name__=='__main__':
    parser = argparse.ArgumentParser(prog='procks', description="Transparent SOCKS5/SOCKS4/HTTP_CONNECT Proxy")
    parser.add_argument('-n', '--nat', action='store_true', help="set bind address to 0.0.0.0 to make pr0cks work from a netfilter FORWARD rule instead of OUTPUT")
    parser.add_argument('--enjoysoft', action='store_true', help="set bind address to 192.168.0.31")
    parser.add_argument('-v', '--verbose', action="store_true", help="print all the connections requested through the proxy")
    parser.add_argument('-c', '--no-cache', action="store_true", help="don't cache dns requests")
    parser.add_argument('--dns-port', default=1053, type=int, help="dns port to listen on (default 1053)")
    parser.add_argument('--dns-server', default="8.8.8.8:53", help="ip:port of the DNS server to forward all DNS requests to using TCP through the proxy (default 8.8.8.8:53)")#208.67.222.222:53
    args=parser.parse_args()

    bind_address="127.0.0.1"
    if args.nat:
        bind_address="0.0.0.0"
    if args.enjoysoft:
        bind_address="192.168.0.31"
    if dnslib_imported:
        try:
            dns_srv, dns_port=args.dns_server.split(':',1)
            dns_port=int(dns_port)
        except Exception as e:
            display("[-] %s"%e)
            display("[-] Invalid dns server : %s"%args.dns_server)
            exit(1)
        resolver = ProxyResolver(dns_srv,dns_port)
        handler = PassthroughDNSHandler # if args.passthrough else DNSHandler
        logger = DNSLogger("request,reply,truncated,error", False)
        udp_server = DNSServer(resolver,
                               port=args.dns_port,
                               address=bind_address,
                               logger=logger,
                               handler=handler)
        udp_server.start_thread()
        display("[+] DNS server started on %s:%s forwarding all DNS trafic to %s:%s using TCP"%(bind_address, args.dns_port, dns_srv, dns_port))

    except KeyboardInterrupt:
        sys.stdout.write("\n")
        sys.exit(0)
    except Exception as e:
        sys.stderr.write(traceback.format_exc())
        sys.exit(1)