#!/usr/bin/env python
# -*- coding: iso-8859-1 -*-

# This is jcbt, a BitTorrent toolkit.
# For more information, see http://www.lysator.liu.se/~jc/jcbt.html

# Copyright (C) 2005  Jörgen Cederlöf <jc@lysator.liu.se>
# 
# This program 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 2
# of the License, or (at your option) any later version.
# 
# This program 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 this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA
# 02111-1307, USA.

VERSION = "0.5"

import sys
import pprint
from BitTornado.bencode import bencode, bdecode
from binascii import b2a_hex

def l2h(l):
    prefixes = ["B", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB"]
    s, i = [], 0
    while l:
        if l%1024:
            if i < len(prefixes):
                p = prefixes[i]
            else:
                p = "*1024^%d" % i
            s.append("%d %s" % (l%1024, p))
        l //= 1024
        i  += 1
    return " + ".join(s[::-1])

def l2hn(l):
    i = 1
    while l > i:
        i *= 2
    return "%s - ( %s )" % (l2h(i), l2h(i-l) or 0)

def prettyl(l):
    return "%d bytes == 0x%X bytes == %s == %s" % (l, l, l2h(l), l2hn(l))

def findranges(l):
    """Return a string describing the list of integers l in a compact
       format. Assumes l is sorted and has no duplicates."""
    if not l:
        return ""
    begin = l[0]
    n = begin
    ranges = []
    for i in l[1:] + [None]:
        if i == n+1:
            n = i
        elif n == begin:
            ranges.append("%d" % n)
            begin = n = i
        else:
            ranges.append("%d-%d" % (begin, n))
            begin = n = i
    
    return ' '.join(ranges)

def getzerosduplicates(pieces, piecesize):
    import sha
    zerodigest = sha.sha("\0"*piecesize).hexdigest()
    pieces = [pieces[x:x+20] for x in xrange(0, len(pieces), 20)]
    zeros = []
    duplicatedigests = {}
    for n, piece in enumerate(pieces):
        digest = b2a_hex(piece)
        if digest == zerodigest:
            zeros.append(n)
        elif digest not in duplicatedigests:
            duplicatedigests[digest] = [n]
        else:
            duplicatedigests[digest].append(n)
    for digest in list(duplicatedigests):
        if len(duplicatedigests[digest]) < 2:
            del duplicatedigests[digest]
    return zeros, duplicatedigests

def getinfohash(torrent, binary=False):
    import sha
    # Dictionaries are always sorted as raw strings when bencoding, so
    # the info hash should be safe to create from the bdecoded
    # dictionary.
    return sha.sha(bencode(torrent['info'])).hexdigest()

def printall(torrent):
    import time
    info_hash = getinfohash(torrent)
    pieces = torrent['info']['pieces']
    torrent['info']['pieces'] = '*** A %d bytes long concatenation of %d hashes ***' \
                                % (len(pieces), len(pieces)/20)
    pretty = pprint.pformat(torrent)
    print "The whole torrent file decoded, with pieces censored:"
    print "    " + pretty.replace('\n', '\n    ')
    if 'length' in torrent['info']:
        totlen = torrent['info']['length']
    else:
        totlen = sum([f['length'] for f in torrent['info']['files']])
    print "Total size: " + prettyl(totlen)
    piecelen = torrent['info']['piece length']
    print "Piece size: " + prettyl(piecelen)
    try:
        print "Creation date: " + time.ctime(torrent['creation date'])
    except KeyError:
        print "No creation date is available in this .torrent."
    zeros, duplicates = getzerosduplicates(pieces, piecelen)
    if zeros:
        print "%d/%d" % (len(zeros), len(pieces)/20),
        print "zero pieces were found: " + findranges(zeros)
    else:
        print "No zero pieces."
    if duplicates:
        print "Duplicate pieces:"
        total = 0
        for dup in duplicates:
            print "  %s: %s" % (dup, findranges(duplicates[dup]))
            total += len(duplicates[dup])
        if total > 0:
            print " Totally %d pieces duplicated over %d pieces." % (len(duplicates), total)
    else:
        print "No duplicate pieces."
    print "Info hash:", info_hash


        
def getsock(beginport=4000):
    import socket
    host = "127.0.0.3"
    port = beginport
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind((host, port))
    return host, port, s

def listen(s):
    s.listen(1)
    conn, addr = s.accept()
    return conn, addr

def server():
    import cgi
    host, port, s = getsock()
    tracker_url = "http://%s:%d/foo" % (host, port)
    peers = []
    responsedict = {'interval': 60,
                    'peers':    peers}
    ids = {}
    while True:
        print
        if peers:
            print "Known peers:"
            for peer in peers:
                print "   %15s:%-5d %s" % (peer['ip'], peer['port'], peer['peer id'])
        else:
            print "No known peers."
        print "Tracker URL: ", tracker_url
        print
        conn, addr = listen(s)
        query = conn.recv(100000)
        print 'Connected by', addr
        for line in query.split('\n'):
            print "remote> ", line
        q = query.split(' ', 2)[1].split('?', 1)[1]
        q = cgi.parse_qs(q)
        peer_id   = q['peer_id'][0]
        peer_port = int(q['port'][0])
        peer_ip   = q.get('ip', [addr[0]])[0]
        print "peer_id, port: %s, %d" % (repr(peer_id), peer_port)
        if peer_id not in ids:
            peers.append({'peer id': peer_id,
                          'ip':      peer_ip,
                          'port':    peer_port})
            ids[peer_id] = None
        response = bencode(responsedict)
        print "me> ", response
        conn.send(response)
        conn.close()

def talktotracker(argv, cmd="jcbt track"):
    def usage():
        print "%s trackerURL info_hash" % cmd
        print "  or"
        print "%s torrentfile" % cmd
        sys.exit(1)
    if   len(argv) == 1:
        torrent = bdecode(open(argv[0]).read())
        tracker = torrent['announce']
        infohash = getinfohash(torrent)
    elif len(argv) == 2:
        tracker, infohash = argv
    else:
        usage()

    for n in range(len(infohash)//2)[::-1]:
        infohash = infohash[:n*2] + "%" + infohash[n*2:]

    import random
    import time
    peer_id = "IGNORE_%s_" % time.strftime('%H:%M:%S')  \
              + "".join([chr(random.randint(ord('a'), ord('z'))) for n in range(4)])

    # Even without event, at least TPB returns me on later queries. 
    # With port 0. (TPB/Hypercube seems to take integer ports mod
    # 2**32 and translate other ports to 0. Or rather stop reading
    # when not digit.)
    url = "%s?info_hash=%s&peer_id=%s&port=0&uploaded=0&downloaded=0&left=1&event=" \
          % (tracker, infohash, peer_id)

    #url = tracker + '?info_hash=%92%0B%D2%0C%97%85%40%B0%C6%A9c%8DjB%AA%92%40%23%15%26&peer_id=T03A-----cpCWNVCoT8w&port=6931&uploaded=0&downloaded=0&left=1&no_peer_id=1&compact=1&key=C79fRm'

    print "Fetching %s" % url

    import re
    match = re.compile('http://([^/:]*)(:[0-9]+)?(/.*)').search(url)
    host, port, path = match.groups()
    if port is None: port = 80
    else:            port = int(port[1:])

    import socket
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    print "Connecting to %s:%d" % (host, port)
    s.connect((host, port))
    query = "GET %s HTTP/1.0\r\n" % path
    print "me> %s" % repr(query)
    s.send(query)
    query = "Accept-Encoding: identity\r\nHost: %s:%s\r\n\r\n" % (host, port)
    print "me> %s" % repr(query)
    s.send(query)
    response, resp = "", s.recv(2**16)
    while resp:
        response = response + resp
        resp = s.recv(2**16)
    lines = response.split('\n')
    s.close()
    for n, line in enumerate(lines):
        print "remote>", repr(line)
        if not line.strip():
            break
    try:
        decoded = bdecode('\n'.join(lines[n+1:]))
        print "remote sent a bencoded dictionary:"
        for line in pprint.pformat(decoded).split('\n'):
            print "   ", line
    except ValueError:
        for line in lines[n+1:]:
            print "remote>", repr(line)
    
def btdiff(torrent0, torrent1):
    size0   = torrent0['info']['piece length']
    size1   = torrent1['info']['piece length']
    hashes0 = torrent0['info']['pieces']
    hashes1 = torrent1['info']['pieces']
    if size0 == size1:
        print "Piece size is %s in both torrents." % l2h(size0)
    else:
        print "Piece sizes differ: %s vs %s. Aborting." % (size0, size1)
        return
    equal = True
    for i in xrange(min(len(hashes0), len(hashes1))/20):
        if hashes0[i*20:i*20+20] != hashes1[i*20:i*20+20]:
            equal = False
            print "Piece %d differ: %s != %s" \
                  % (i, b2a_hex(hashes0[i*20:i*20+20]),
                        b2a_hex(hashes1[i*20:i*20+20]))
    if equal:
        print "All pieces were identical."
    if len(hashes0) == len(hashes1):
        print "Both torrents contain %d pieces." % len(hashes0)
    else:
        print "%d pieces are only present in the %s file." % \
              (abs(len(hashes0) - len(hashes1))/20,
               len(hashes0)>len(hashes1) and "first" or "second")
        

def main(argv):
    def usage(e=1):
        if len(argv) > 0:
            me = argv[0]
        else:
            me = "jcbt"
        print "This is jcbt version %s." % VERSION
        print "Copyright (C) 2005 Jörgen Cederlöf <jc@lysator.liu.se>"
        print "Usage:"
        print "%s help" % me
        print "    Show this help."
        print "%s print file.torrent" % me
        print "    Shows all information contained in file.torrent, except hashes."
        print "    Unlike other programs with similar purposes, even all unknown"
        print "    fields are shown."
        print "    Decodes total size and piece size"
        print "    Detects duplicate and zero pieces"
        print "%s diff first.torrent second.torrent" % me
        print "    Find differences between the hashes in two torrent files."
        print "%s track trackerURL info_hash" % me
        print "%s track file.torrent" % me
        print "    Connect to the tracker and receive information about other clients."
        print "%s server" % me
        print "    Start a simple tracker connecting all clients calling it."
        print "    Very useful for syncing two clients downloading the same files from"
        print "    two different incompatible torrents using a third client acting as"
        print "    a proxy. With this server the proxy client can connect to the real"
        print "    clients without relying on the external tracker."
        sys.exit(e)
    try:
        op      = sys.argv[1]
    except:
        usage()
    if   op in ('help', '--help'):
        usage(0)
    if   op == 'print':
        torrent = bdecode(open(sys.argv[2]).read())
        printall(torrent)
    elif op == 'server':
        server()
    elif op == 'track':
        talktotracker(sys.argv[2:], " ".join(sys.argv[:2]))
    elif op == 'diff':
        torrent0 = bdecode(open(sys.argv[2]).read())
        torrent1 = bdecode(open(sys.argv[3]).read())
        btdiff(torrent0, torrent1)
    else:
        usage()

if __name__ == "__main__":
    main(sys.argv)