HttpUploadComponent/server.py

import yaml
import sys
import shutil
import signal
import logging
import string
import hashlib
import random
import os
from threading import Thread
from threading import Lock
try:
    # Python 3
    from http.server import HTTPServer, BaseHTTPRequestHandler
    from socketserver import ThreadingMixIn
except ImportError:
    # Python 2
    from BaseHTTPServer import HTTPServer, BaseHTTPRequestHandler
    from SocketServer import ThreadingMixIn
import sleekxmpp
from sleekxmpp.componentxmpp import ComponentXMPP

global files
global files_lock
global config

class MissingComponent(ComponentXMPP):
    def __init__(self, jid, secret):
        ComponentXMPP.__init__(self, jid, secret, "localhost", 5347)
        self.register_plugin('xep_0030')
        self.register_plugin('upload',module='plugins.upload')
        self.add_event_handler('request_upload_slot',self.request_upload_slot)

    def request_upload_slot(self, iq):
        global config
        global files
        global files_lock
        request = iq['request']
        maxfilesize = int(config['max_file_size'])
        if not request['filename'] or not request['size']:
            self._sendError(iq,'modify','bad-request','please specify filename and size')
        elif maxfilesize < int(request['size']):
            self._sendError(iq,'modify','not-acceptable','file too large. max file size is '+str(maxfilesize))
        elif 'whitelist' not in config or iq['from'].domain in config['whitelist']:
            sender = iq['from'].bare
            sender_hash = hashlib.sha1(sender.encode()).hexdigest()
            filename = request['filename']
            folder = ''.join(random.SystemRandom().choice(string.ascii_uppercase + string.digits) for _ in range(len(sender_hash)))
            sane_filename = "".join([c for c in filename if c.isalpha() or c.isdigit() or c=="."]).rstrip()
            path = sender_hash+'/'+folder
            if sane_filename:
                path += '/'+sane_filename
            with files_lock:
                files.add(path)
            print(path)
            reply = iq.reply()
            reply['slot']['get'] = config['get_url'] + '/' + path
            reply['slot']['put'] = config['put_url'] + '/' + path
            reply.send()
        else:
           self. _sendError(iq,'cancel','not-allowed','not allowed to request upload slots')

    def _sendError(self, iq, error_type, condition, text):
        reply = iq.reply()
        iq.error()
        iq['error']['type'] = error_type
        iq['error']['condition'] = condition
        iq['error']['text'] = text
        iq.send()

class HttpHandler(BaseHTTPRequestHandler):
    def do_PUT(self):
        print('do put')
        global files
        global files_lock
        global config
        path = self.path[1:]
        length = int(self.headers['Content-Length'])
        maxfilesize = int(config['max_file_size'])
        if maxfilesize < length:
            self.send_response(400,'file too large')
            self.end_headers()
        else:
            print('path: '+path)
            files_lock.acquire()
            if path in files:
                files.remove(path)
                files_lock.release()
                filename = config['storage_path'] + path
                os.makedirs(os.path.dirname(filename))
                remaining = length
                f = open(filename,'wb')
                data = self.rfile.read(4096)
                while data and remaining >= 0:
                    remaining -= len(data)
                    f.write(data)
                    data = self.rfile.read(min(4096,remaining))
                f.close()
                self.send_response(200,'ok')
                self.end_headers()
            else:
                files_lock.release()
                self.send_response(403,'invalid slot')
                self.end_headers()

    def do_GET(self):
        global config
        path = self.path[1:].replace('../','').replace('./','')
        slashcount = path.count('/')
        if slashcount < 1 or slashcount > 2:
            self.send_response(404,'file not found')
            self.end_headers()
        else:
            filename = config['storage_path']+'/'+path
            print('requesting file: '+filename)
            try:
                with open(filename,'rb') as f:
                    self.send_response(200)
                    self.send_header("Content-Type", 'application/octet-stream')
                    self.send_header("Content-Disposition", 'attachment; filename="{}"'.format(os.path.basename(filename)))
                    fs = os.fstat(f.fileno())
                    self.send_header("Content-Length", str(fs.st_size))
                    self.end_headers()
                    shutil.copyfileobj(f, self.wfile)
            except FileNotFoundError:
                self.send_response(404,'file not found')
                self.end_headers()

    def do_HEAD(self):
        global config
        path = self.path[1:].replace('../','').replace('./','')
        slashcount = path.count('/')
        if slashcount < 1 or slashcount > 2:
            self.send_response(404,'file not found')
            self.end_headers()
        else:
            try:
                filename = config['storage_path']+'/'+path
                self.send_response(200,'OK')
                self.send_header("Content-Type", 'application/octet-stream')
                self.send_header("Content-Disposition", 'attachment; filename="{}"'.format(os.path.basename(filename)))
                self.send_header("Content-Length",str(os.path.getsize(filename)))
                self.end_headers()
            except FileNotFoundError:
                self.send_response(404,'file not found')
                self.end_headers()


class ThreadedHTTPServer(ThreadingMixIn, HTTPServer):
    """Handle requests in a separate thread."""

if __name__ == "__main__":
    global files
    global files_lock
    global config

    with open('config.yml','r') as ymlfile:
        config = yaml.load(ymlfile)

    files = set()
    files_lock = Lock()
    logging.basicConfig(level=logging.DEBUG,
                            format='%(levelname)-8s %(message)s')
    server = ThreadedHTTPServer(('0.0.0.0', config['http_port']), HttpHandler)
    xmpp = MissingComponent(config['jid'],config['secret'])
    if xmpp.connect():
        xmpp.process(block=False)
        print("connected")
        server.serve_forever()
    else:
        print("unable to connect")