from bencode import bencode, bdecode
import socket
from threading import Thread
import random
from hashlib import sha1
import time
import json
import struct
from Queue import Queue
import binascii
from SimpleXMLRPCServer import SimpleXMLRPCServer

BOOTSTRAP_NODES = [
    ("router.bittorrent.com", 6881),
    ("dht.transmissionbt.com", 6881),
    ("router.utorrent.com", 6881)
] 

q = Queue()

def entropy(bytes):
    s = ""
    for i in range(bytes):
        s += chr(random.randint(0, 255))
    return s


class DHT(Thread):
    def __init__(self):
        Thread.__init__(self)
        self.ufd = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
        self.ufd.bind(("0.0.0.0", 6881))
        self.nid = self.rand_id()
        self.next_sequence = 0
        self.version = 'CWxx'
        self.new_nodes = []
        self.db = open("data.txt", "a")

    def rand_id(self):
        hash = sha1()
        hash.update( entropy(20) )
        return hash.digest()
    
    def new_sequence(self):
        seq = '%8d' % (self.next_sequence % 65536*65535)
        self.next_sequence += 1
        return seq

    def parse_nodes(self, nodes):
        n = []
        length = len(nodes)
        if (length % 26) != 0: 
            return n
        for i in range(0, length, 26):
            nid = nodes[i:i+20]
            ip = socket.inet_ntoa(nodes[i+20:i+24])
            port = struct.unpack("!H", nodes[i+24:i+26])[0]
            n.append( (ip, port) )
        return list(set(n))

    def run(self):
        while True:
            try:
                data, addr = self.ufd.recvfrom(65536)
            except:
                print "ERROR:recv data"
                continue
            try:
                msg = bdecode(data)
            except:
                print "bdecode error"
                continue
            self.recieve_msg(msg, addr)

    def recieve_msg(self, msg, addr):
        #print "++++++++++++RESPONSE MSG++++++++++++++"
        #print msg
        #print "++++++++++++RESPONSE MSG++++++++++++++"
        try:
            t = msg['y']
        except:
            return
        if t == "q":
            self.parse_request(msg, addr)
        elif t == "r":
            self.parse_response(msg, addr)
        else:
            print "IT'S AN ERROR MSG"

    def send_message(self, addr, msg):
        data = bencode(msg)
        try:
            self.ufd.sendto(data, addr)
        except:
            return
            #print '[ERROR] sending', msg, 'to', addr

    def parse_request(self, msg, addr):
        query_method = msg['q']
        resp_msg = {}
        if query_method == "ping":
            print "++++++ GOT A PING REQUEST++++++++"
            resp_msg = {
                "t": self.new_sequence(),
                "y": "r",
                "r": {"id": self.nid}
            }
        elif query_method == "find_node":
            print "++++++ GOT A FIND_NODE REQUEST++++++++"
            resp_msg = {
                "t": msg['t'],
                "y": "r",
                "r": { "id":self.nid, "nodes":""}
            }
        elif query_method == "get_peers":
            print "++++++ GOT A GET_PEERS REQUEST++++++++"
            print "----------ADDR:", addr
            resp_msg = {
                "id": self.nid,
                "t": msg["t"],
                "y": "r",
                #"r": {"id":self.nid, "token":"tokenstr", "values":""}
                "r": {"id":self.nid, "token":"tokenstr", "nodes":""}
            }
            self.log_infohash(msg, addr)
        elif query_method == "announce_peer":
            print "++++++ GOT A ANNOUNCE_PEER REQUEST++++++++"
        # send reply
        self.send_message(resp_msg, addr)

    def log_infohash(self, msg, addr):
        infohash = binascii.hexlify(msg["a"]["info_hash"])
        print ">>>>>>>>>>>>>  [%s] <<<<<<<<<<<<<<<<<<<" % infohash
        q.put(infohash)
        #self.db.write(infohash + "\n")
        #self.db.flush()

    def parse_response(self, msg, addr):
        try:
            l = self.parse_nodes(msg['r']['nodes'])
        except:
            return
        for node in l:
            self.new_nodes.append(node)

    def ping(self, addr):
        m = {
            'v': self.version,
            't': self.new_sequence(),
            'y': 'q',
            'q': 'ping',
            'a': {
                'id': self.nid
            }
        }
        self.send_message(addr, m)

    def find_node(self, addr, target_id):
        m = {
            'v': self.version,
            't': self.new_sequence(),
            'y': 'q',
            'q': 'find_node',
            'a': {
                'id': self.nid,
                'target': target_id,
            }
        }
        self.send_message(addr, m)

    def get_peers(self, addr, info_hash):
        m = {
            'v': self.version,
            't': self.new_sequence(),
            'y': 'q',
            'q': 'get_peers',
            'a': {
                'id': self.nid,
                'info_hash': info_hash,
            }
        }
        self.send_message(addr, m)

    def distance(self, a, b):
        dist = sum([bin(ord(x) ^ ord(y)).count('1') for x,y in zip(a,b)])
        return dist

    def boot(self):
        print "+++++++++++++BOOT START+++++++++++++++"
        for addr in BOOTSTRAP_NODES:
            self.find_node(addr, self.rand_id())
    
    def auto_find(self):
        print "+++++++++++++AUTO FIND+++++++++++++++"
        for addr in self.new_nodes:
            self.get_peers(addr, self.rand_id())

def get_hash():
    return q.get()

def rpc_server():
    server = SimpleXMLRPCServer(('0.0.0.0', 1989))
    server.register_function(get_hash, 'get_hash')
    server.serve_forever()


if __name__ == '__main__':
    rpc_thread = Thread(target=rpc_server)
    rpc_thread.start()
    d = DHT()
    d.start()
    d.boot()
    while True:
        d.auto_find()
        time.sleep(1)
    d.join()