Share files through LAN (Broadcast)

co.py

#!/usr/bin/env python
import SocketServer
import struct
import shutil
import hashlib
import zlib
from StringIO import StringIO
import os
try:
    import wx
except:
    wx = None

# rsync from http://blog.liw.fi/posts/rsync-in-python/

block_size = 4096

def signature(f):
    while True:
        block_data = f.read(block_size)
        if not block_data:
            break
        yield (zlib.adler32(block_data), hashlib.md5(block_data).digest())

class RsyncLookupTable(object):

    def __init__(self, checksums):
        self.dict = {}
        for block_number, c in enumerate(checksums):
            weak, strong = c
            if weak not in self.dict:
                self.dict[weak] = dict()
            self.dict[weak][strong] = block_number

    def __getitem__(self, block_data):
        weak = zlib.adler32(block_data)
        subdict = self.dict.get(weak)
        if subdict:
            strong = hashlib.md5(block_data).digest()
            return subdict.get(strong)
        return None

def delta(sigs, f):
    table = RsyncLookupTable(sigs)
    block_data = f.read(block_size)
    while block_data:
        block_number = table[block_data]
        if block_number:
            yield (block_number * block_size, len(block_data))
            block_data = f.read(block_size)
        else:
            yield block_data[0]
            block_data = block_data[1:] + f.read(1)

def patch(outputf, deltas, old_file):
    for x in deltas:
        if type(x) == str:
            outputf.write(x)
        else:
            offset, length = x
            old_file.seek(offset)
            outputf.write(old_file.read(length))
            



CHKSUM = "<i16c"

def netstr(*args):
    return ''.join("%d:%s" % (len(s), s) for s in args)

def getstr(sock):
     size = 0
     c = "0"
     while c != ":" and c.isdigit():
        size = size * 10 + int(c)
        c = sock.recv(1)
     if size == 0:
        return ""
     return sock.recv(size)
     
def strchecksums(s):
    for idx in xrange(0, len(s), struct.calcsize(CHKSUM)):
        chk = struct.unpack(CHKSUM, s[idx:idx + struct.calcsize(CHKSUM)])
        yield chk[0], ''.join(chk[1:]) 

def chksumstr(sig):
    a, b = sig
    return struct.pack(CHKSUM, a, *b)



def patch_file_from(old, addr):
    old = old.split("/")[-1]
    netname = netstr(old)
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect((addr, 7331))
    try:
        oldf = open(old, "rb")
    except IOError:
        oldf = StringIO()
    sigs = ''.join(chksumstr(s) for s in signature(oldf))
    netsigs = netstr(sigs)
    sock.sendall(netname + netsigs)
    newf = StringIO()
    
    diff = []
    for i in xrange(int(getstr(sock))):
        t = getstr(sock)
        if t == "LITERAL": 
            diff.append(getstr(sock))
        else:
            d0 = int(getstr(sock))
            d1 = int(getstr(sock))
            diff.append((d0, d1))
    patch(newf, diff, oldf)
    oldf.close()
    oldf = open(old, "wb")
    newf.seek(0)
    shutil.copyfileobj(newf, oldf)
    oldf.close()
    newf.close()


class RsyncRequestHandler(SocketServer.BaseRequestHandler):

    def handle(self):
        s = self.request
        filename = getstr(s)
        checksums = list(strchecksums(getstr(s)))
        try:
            f = open(filename, "rb")
        except IOError:
            print "no file"
            f = StringIO()
        diff = list(delta(checksums, f))
        s.sendall(netstr(str(len(diff))))
        for d in diff:
            if isinstance(d, tuple):
                s.sendall(netstr("TUPLE", str(d[0]), str(d[1])))
            else:
                s.sendall(netstr("LITERAL", d))
        f.close()
        return


class RsyncServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer):
    allow_reuse_address = True

    def __init__(self, address=('', 7331), filehandler=None):
        SocketServer.TCPServer.__init__(self, address, RsyncRequestHandler)
        self.filehandler = filehandler


class DiscoveryRequestHandler(SocketServer.BaseRequestHandler):

    def handle(self):
        print "New discovery request"
        data, client = self.request
        if data[0] == "-":
            try:
                os.unlink(data[1:])
                return
            except:
                pass
        elif data[0] == "+":
            patch_file_from(data[1:], self.client_address[0])
        else:
            if (hasattr(self.server, "filehandler") and 
                    callable(self.server.filehandler)):
                self.server.filehandler(self.client_address[0], data) 


class DiscoveryServer(SocketServer.ThreadingMixIn, SocketServer.UDPServer):
    allow_reuse_address = True 
    
    def __init__(self, filehandler=None, address=('', 1337)):
        SocketServer.UDPServer.__init__(self, address, DiscoveryRequestHandler)
        self.filehandler = filehandler
        
def bcast_file(f, change="+", discovery_bcast_addr=('255.255.255.255', 1337)):    
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    s.bind(("", 1337))
    s.sendto(change + f, discovery_bcast_addr)
    
   
if wx:   
    class FileDropTarget(wx.FileDropTarget):
    
       def __init__(self, window):
            wx.FileDropTarget.__init__(self)
            self.window = window

       def OnDropFiles(self, x, y, filenames):
            bcast_file('|'.join(filenames), change="*")
            for name in filenames:
                bcast_file(name)


    class MainWindow(wx.Frame):
        def __init__(self, parent, id, title):
            wx.Frame.__init__(self, parent, id, title, size=(300, 300))  
            dt = FileDropTarget(self)
            self.SetDropTarget(dt)
            self.Centre()
            self.Show(True)
     

if __name__ == '__main__':
    import socket
    import threading

    discovery_server = DiscoveryServer()

    discovery_thread = threading.Thread(target=discovery_server.serve_forever)
    discovery_thread.setDaemon(True) # don't hang on exit
    discovery_thread.start()
    
    rsync_server = RsyncServer()

    rsync_thread = threading.Thread(target=rsync_server.serve_forever)
    rsync_thread.setDaemon(True) # don't hang on exit
    rsync_thread.start()
    if wx:
        app = wx.App()
        w = MainWindow(None, -1, '<title goes here>')
        
        app.MainLoop()
    else:
        f = raw_input("File >")
        while f != "":
            bcast_file(f)
            f = raw_input("File >")