HttpUploadComponent/server.py

198 lines
7.0 KiB
Python
Raw Normal View History

2015-07-04 15:41:45 -04:00
#!/usr/bin/env python
import argparse
import errno
2015-06-28 05:21:30 -04:00
import hashlib
import logging
import mimetypes
2015-06-28 05:21:30 -04:00
import os
import random
import shutil
import signal
import sleekxmpp
2015-07-10 08:28:59 -04:00
import ssl
import string
import sys
import yaml
from sleekxmpp.componentxmpp import ComponentXMPP
2015-06-28 05:21:30 -04:00
from threading import Lock
from threading import Thread
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
try:
FileNotFoundError
except NameError:
# Python 2
class FileNotFoundError(IOError):
def __init__(self, message=None, *args):
super(FileNotFoundError, self).__init__(args)
self.message = message
self.errno = errno.ENOENT
def __str__(self):
return self.message or os.strerror(self.errno)
2015-06-28 05:21:30 -04:00
LOGLEVEL=logging.DEBUG
2015-06-28 05:21:30 -04:00
global files
global files_lock
global config
def normalize_path(path):
"""
Normalizes the URL to prevent users from grabbing arbitrary files via `../'
and the like.
"""
return os.path.normcase(os.path.normpath(path))
2015-06-28 05:21:30 -04:00
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 = os.path.join(sender_hash, folder)
2015-06-28 05:21:30 -04:00
if sane_filename:
path = os.path.join(path, sane_filename)
2015-06-28 05:21:30 -04:00
with files_lock:
files.add(path)
print(path)
reply = iq.reply()
reply['slot']['get'] = os.path.join(config['get_url'], path)
reply['slot']['put'] = os.path.join(config['put_url'], path)
2015-06-28 05:21:30 -04:00
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 = normalize_path(self.path[1:])
2015-06-28 05:21:30 -04:00
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 = os.path.join(config['storage_path'], path)
2015-06-28 05:21:30 -04:00
os.makedirs(os.path.dirname(filename))
remaining = length
2015-07-22 23:21:28 -04:00
with open(filename,'wb') as f:
data = self.rfile.read(4096)
while data and remaining >= 0:
remaining -= len(data)
f.write(data)
data = self.rfile.read(min(4096,remaining))
2015-06-28 05:21:30 -04:00
self.send_response(200,'ok')
self.end_headers()
else:
files_lock.release()
self.send_response(403,'invalid slot')
self.end_headers()
2015-07-23 01:42:24 -04:00
def do_GET(self, body=True):
2015-06-28 05:21:30 -04:00
global config
path = normalize_path(self.path[1:])
2015-06-28 05:21:30 -04:00
slashcount = path.count('/')
if path[0] in ('/', '\\') or slashcount < 1 or slashcount > 2:
2015-06-28 05:21:30 -04:00
self.send_response(404,'file not found')
self.end_headers()
else:
filename = os.path.join(config['storage_path'], path)
2015-06-28 05:21:30 -04:00
print('requesting file: '+filename)
try:
with open(filename,'rb') as f:
self.send_response(200)
mime, _ = mimetypes.guess_type(filename)
if mime is None:
mime = 'application/octet-stream'
self.send_header("Content-Type", mime)
2015-06-28 05:21:30 -04:00
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()
2015-07-23 01:42:24 -04:00
if body:
shutil.copyfileobj(f, self.wfile)
2015-06-28 05:21:30 -04:00
except FileNotFoundError:
self.send_response(404,'file not found')
self.end_headers()
def do_HEAD(self):
2015-07-23 01:42:24 -04:00
self.do_GET(body=False)
2015-06-28 05:21:30 -04:00
class ThreadedHTTPServer(ThreadingMixIn, HTTPServer):
"""Handle requests in a separate thread."""
2015-06-28 05:21:30 -04:00
if __name__ == "__main__":
global files
global files_lock
global config
parser = argparse.ArgumentParser()
parser.add_argument("-c", "--config", type=unicode, default='config.yml', help='Specify alternate config file.')
parser.add_argument("-l", "--logfile", type=unicode, default=None, help='File where the server log will be stored. If not specified log to stdout.')
args = parser.parse_args()
with open(args.config,'r') as ymlfile:
2015-06-28 05:21:30 -04:00
config = yaml.load(ymlfile)
files = set()
files_lock = Lock()
logging.basicConfig(level=LOGLEVEL,
format='%(asctime)-24s %(levelname)-8s %(message)s',
filename=args.logfile)
2015-07-22 02:15:39 -04:00
server = ThreadedHTTPServer((config['http_address'], config['http_port']), HttpHandler)
2015-07-10 08:28:59 -04:00
if 'keyfile' in config and 'certfile' in config:
server.socket = ssl.wrap_socket(server.socket, keyfile=config['keyfile'], certfile=config['certfile'])
2015-06-28 05:21:30 -04:00
xmpp = MissingComponent(config['jid'],config['secret'])
if xmpp.connect():
xmpp.process()
2015-06-28 05:21:30 -04:00
print("connected")
server.serve_forever()
else:
print("unable to connect")