#!/usr/bin/python
#
# tinyswarm.py -- P2P swarming in <150 lines(*)

# Copyright (c) 2004, Brian St. Pierre <brian@bstpierre.org>

# This will swarm a file using a small, known network of peers. No
# attempt at peer discovery is made. Inspired by tinyp2p, but 10x
# longer -- and 10x easier to understand. Only one file may be swarmed
# over this network at a time.

# Probably works with any python2. Tested on 2.2.[123], 2.3.3.

# Usage:
# tinyswarm.py <filename> <peer-ip-address> [<peer-ip-address> ...]
#
# If 'filename' is '-', we will connect to each peer and start
# swarming from whoever has pieces (hunks) of the file. Otherwise, we will
# offer hunks to anyone that connects to us.
#
# You'll need to have port 8877 open on your firewall, or adjust
# LISTEN_PORT below.

# THIS IS INCREDIBLY INSECURE. If someone knows your peer address,
# they may be able to force you to swarm a file.

# This code is not thread-safe, which causes a bug about 1 in 10 times
# (in my testing -- YMMV) when one of the peers nears completion of
# receiving the file, and occasionally when starting up.

# Use this program however you want, but THERE IS ABSOLUTELY NO
# WARRANTY for this program. If you break it, you get to keep both
# pieces.

# (*) 150 lines doesn't count comments, whitespace, or docstrings.

import os
import random
import socket
import sys
import threading
import time

LISTEN_PORT = 8878
TINYSWARM_VERSION = 0.1

class SwarmFile:
    '''A file to be swarmed. Hunks are downloaded and then made
    available to other peers.'''

    _HUNK_SIZE = 1400

    def __init__(self, filename):
        self._name = ''
        self._hunks = []
        if filename != '-':
            self._name = os.path.basename(filename)
            self._hunkify(open(filename))
        return

    def _hunkify(self, f):
        '''Split the file into roughly equal sized hunks. (The last
        hunk may be truncated.)'''
        hunk = f.read(self._HUNK_SIZE)
        while hunk != '':
            self._hunks.append(hunk)
            hunk = f.read(self._HUNK_SIZE)
        return

    def name(self):
        '''Get the name of the file.'''
        return self._name

    def set_name(self, name):
        '''Sets the name of the file.'''
        assert(name != '')
        self._name = name
        return

    def num_hunks(self):
        '''Return the number of hunks.'''
        return len(self._hunks)

    def alloc_hunks(self, count):
        '''Create space for all of the hunks.'''
        if len(self._hunks) == 0:
            self._hunks = [None] * count
        else:
            assert(len(self._hunks) == count), (len(self._hunks), count)
        return

    def empty_hunks(self):
        '''Return the number of missing hunks.'''
        return self._hunks.count(None)

    def complete(self):
        '''Return true if we are not missing any hunks.'''
        return self.empty_hunks() == 0

    def set_hunk(self, n, hunk):
        '''Set the given hunk.'''
        if self._hunks[n] == None:
            self._hunks[n] = hunk
        else:
            assert(self._hunks[n] == hunk)
        return

    def fetch_hunk(self, peer):
        '''Ask the peer for a random hunk that we do not yet have.'''
        n = random.randint(0, self.num_hunks())
        while n < self.num_hunks():
            if self._hunks[n] == None:
                peer.send('GET %s\n' % n)
                resp = peer.recv(1024)
                # If not 'OK', the peer doesn't have this hunk,
                # try the next one.
                if resp.find('OK') == 0:
                    ok, size = resp.strip().split(' ')
                    size = int(size)
                    self.set_hunk(n, peer.recv(size))
                    break
            n += 1
        return

    def get_hunk(self, n):
        '''Retrieve the given hunk.'''
        return self._hunks[n]

    def write(self):
        '''File must be complete. Write the file using the given
        filename. Will not overwrite if the file has been written
        already. Refuses to write files with slash in the name.'''
        assert(self.complete())
        assert(self._name.find('/') == -1) ## crude safety
        if os.path.exists(self._name):
            return os.path.getsize(self._name)
        f = open(self._name, 'w')
        for h in self._hunks:
            f.write(h)
        f.close()
        return os.path.getsize(self._name)

class SwarmPeer:
    '''A swarming peer. Swarms the given file.'''

    def __init__(self, sfile, ip = ''):
        '''Connect to the given peer. If ip is empty, listen for
        incoming connections. This starts a daemon thread and
        returns.'''
        self._sfile = sfile
        self._ip = ip
        self._port = LISTEN_PORT
        if len(ip) == 0:
            target = self._listen
        else:
            target = self._connect
        thr = threading.Thread(target = target)
        thr.setDaemon(1)
        thr.start()
        return

    def _listen(self):
        '''Wait for other peers to connect. This is a send-only
        connection.'''
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.bind(('', self._port))
        s.listen(5)
        while 1:
            conn, addr = s.accept()
            greeting = 'TINYSWARM %s\n' % TINYSWARM_VERSION
            conn.send(greeting)

            # Wait until we have a file to swarm.
            while self._sfile.name() == '-':
                time.sleep(1)

            conn.send('FILE %s %d\n' % (self._sfile.name(),
                                        self._sfile.num_hunks()))

            # Start a non-daemon thread to send hunks.
            threading.Thread(target = self._deliver,
                             args = [conn]).start()
        return

    def _deliver(self, conn):
        '''Send hunks as they are requested.'''
        req = conn.recv(1024).strip()
        while req != 'BYE':
            if req.find('GET') == 0:
                hunkNum = int(req[4:].strip())
                hunk = self._sfile.get_hunk(hunkNum)
                if hunk != None:
                    conn.send('OK %d\n' % len(hunk))
                    conn.send(hunk)
                else:
                    conn.send('NAK\n')
            req = conn.recv(1024).strip()
        return

    def _connect(self):
        '''Connect to the peer and request hunks. This is a
        receive-only connection.'''
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        while 1:
            try:
                s.connect((self._ip, self._port))
                break
            except socket.error:
                time.sleep(1)
        greeting = 'TINYSWARM %s\n' % TINYSWARM_VERSION
        recv = s.recv(1024)
        assert(recv.strip() == greeting.strip())
        fileCmd, filename, numHunks = s.recv(1024).strip().split(' ')
        self._sfile.set_name(filename)
        self._sfile.alloc_hunks(int(numHunks))
        while not self._sfile.complete():
            self._sfile.fetch_hunk(s)
        self._sfile.write()
        print 'Done receiving file "%s".' % self._sfile.name()
        return
        

if __name__ == '__main__':
    random.seed()

    # In order to keep it short & simple, there's no argument
    # validation.
    sfile = SwarmFile(sys.argv[1])

    # Start up our listener.
    SwarmPeer(sfile, '')

    peers = []
    for peer in sys.argv[2:]:
        peers.append(SwarmPeer(sfile, peer))

    # Peers run as daemon threads. Sit here and hang out until
    # interrupted.
    while 1:
        time.sleep(1000)