simple script to easily share files

share

#!/usr/bin/env python

#pip install pybonjour pyzmq

import argparse
import os
import pybonjour
import select
import sys
import sys
import zmq

service_name = "share_files"
service_port = 4545
TYPE = "_http._tcp"

ip = None
port = None
done = False

debug = False


class ZeroconfServer:

    def __init__(self, name, port, regtype='_http._tcp'):
        self.name = name
        self.port = port
        self.regtype = regtype
        self.published = False

    def register_callback(self, sdRef, flags, errorCode, name, regtype, domain):
        if errorCode == pybonjour.kDNSServiceErr_NoError:
            if debug:
                print 'Registered service:'
                print '  name    =', name
                print '  regtype =', regtype
                print '  domain  =', domain
            self.published = True
        else:
            self.published = False

    def publish(self):
        self.sdRef = pybonjour.DNSServiceRegister(name = self.name,
                regtype = self.regtype,
                port = self.port,
                callBack = self.register_callback)

        while not self.published:
            ready = select.select([self.sdRef], [], [])
            if self.sdRef in ready[0]:
                pybonjour.DNSServiceProcessResult(self.sdRef)

    def unpublish(self):
        if self.published and self.sdRef:
            self.sdRef.close()


class ZeroconfClient:

    def __init__(self, name, port, regtype='_http._tcp', timeout=5):
        self.name = name
        self.port = port
        self.regtype = regtype
        self.timeout = timeout
        self.resolved = False

    def resolve_callback(self, sdRef, flags, interfaceIndex, errorCode, fullname,
                         hosttarget, hostport, txtRecord):
        if errorCode == pybonjour.kDNSServiceErr_NoError:
            if debug:
                print 'Resolved service:'
                print '  fullname   =', fullname
                print '  hosttarget =', hosttarget
                print '  port       =', hostport

            if fullname.startswith(service_name):
                self.resolved = True
                self.hosttarget = hosttarget
                self.hostport = hostport

    def resolve(self):
        resolve_sdRef = pybonjour.DNSServiceResolve(0,
                0,
                self.name,
                self.regtype,
                'local.',
                self.resolve_callback)

        try:
            while not self.resolved:
                ready = select.select([resolve_sdRef], [], [], self.timeout)
                if resolve_sdRef not in ready[0]:
                    if debug:
                        print 'Resolve timed out'
                    break
                pybonjour.DNSServiceProcessResult(resolve_sdRef)
        finally:
            resolve_sdRef.close()


        return self.hosttarget, self.hostport


def get_files():
    global ip
    global port

    context = zmq.Context()
    socket = context.socket(zmq.REQ)

    socket.connect('tcp://%s:%s' % (ip, port))
    # send the desired file to the server
    socket.send('*')

    files = ''
    while True:
        # Start grabing data
        data = socket.recv()
        # Write the chunk to the file
        files += data
        if not socket.getsockopt(zmq.RCVMORE):
            # If there is not more data to send, then break
            break

    return files.splitlines()


def get_file(path):
    '''
A VERY simple client to get files from the zeromq VERY simple file server
'''
    global ip
    global port
    # Open up the file we are going to write to

    directory = os.path.dirname(path)

    if directory:
        if not os.path.exists(directory):
            os.makedirs(directory)

    dest = open(path, 'w+')
    # Set up the zeromq context and socket
    context = zmq.Context()
    socket = context.socket(zmq.REQ)

    socket.connect('tcp://%s:%s' % (ip, port))
    # send the desired file to the server
    socket.send(path)

    while True:
        # Start grabing data
        data = socket.recv()
        # Write the chunk to the file
        dest.write(data)
        if not socket.getsockopt(zmq.RCVMORE):
            # If there is not more data to send, then break
            break


def list_files():
    result = []
    for (dirpath, dirname, filenames) in os.walk('.'):
        for filename in filenames:
            result.append(os.path.join(dirpath, filename))

    result.append('')
    return '\n'.join(result)

def server(files=None):
    '''
Setup a VERY simple file server
'''
    # Set up the zeromq context and REP socket
    context = zmq.Context(1)
    sock = context.socket(zmq.REP)
    sock.bind('tcp://*:%s' % service_port)

    if not files:
        files = list_files()
    # Start the server loop
    while True:
        # Recieve the location of the file to serve
        msg = sock.recv()

        if msg == '*':
            sock.send(files)
            continue

        # Verify that the file is available and in current directory
        if not os.path.isfile(msg) or msg not in files:
            sock.send('')
            continue
        # Open the file for reading
        fn = open(msg, 'rb')
        stream = True
        # Start reading in the file
        while stream:
            # Read the file bit by bit
            stream = fn.read(128)
            if stream:
                # If the stream has more to send then send more
                sock.send(stream, zmq.SNDMORE)
            else:
                # Finish it off
                sock.send(stream)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Share Files')
    parser.add_argument('-d', '--download', action='store_true', help='downloads files from server')
    parser.add_argument('-f', '--files', metavar='filename', type=str, nargs='*')

    args = vars(parser.parse_args())

    if args['download']:
        client = ZeroconfClient(service_name, service_port)
        ip, port = client.resolve()

        if not done:
            for filename in get_files():
                done = True
                print 'downloading file: %s' % filename
                get_file(filename)
            sys.exit()

    service = ZeroconfServer(name=service_name, port=service_port)
    service.publish()

    files = None
    if args['files']:
        files = args['files']
        files.append('')
        files = '\n'.join(files)

    server(files)
    service.unpublish()