import os import io import sys import yaml import string import hashlib from collections import defaultdict from Cryptodome.Cipher import AES from Cryptodome.Util import Padding from aiohttp import web class AttrDict(dict): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.__dict__ = self def update(self, *d, **kwargs): for key, val in (d[0] if d else kwargs).items(): setattr(self, key, val) def __getattr__(self, item): return self.setdefault(item, AttrDict()) def cast_to_ad(d): if not isinstance(d, AttrDict): d = AttrDict(d) for k, v in dict(d.items()).items(): if " " in k: del d[k] d[k.replace(" ", "_")] = v if isinstance(v, dict): d[k] = cast_to_ad(v) return d def sizeof_fmt(num, suffix='B'): for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']: if abs(num) < 1024.0: return f"{num:3.1f}{unit}{suffix}" num /= 1024.0 return f"{num:.1f}Yi{suffix}" async def prepare(_, handler): async def prepare_handler(req): if 'acc' not in req.match_info: return web.Response(text='bad request', status=400) return await handler(req, req.match_info["acc"], file_db[req.match_info["acc"]]) return prepare_handler async def handle_upload(req, acc, acc_db): if conf.auth and req.headers.get('auth') not in conf.auth_tokens: return web.Response(text='access denied', status=403) reader = await req.multipart() file = await reader.next() filename = os.path.basename(file.filename) if not os.path.isdir(f'{conf.data_path}/{acc}'): os.mkdir(f'{conf.data_path}/{acc}') for _ in range(100): hb = os.urandom(conf.url_hash_len//2) h = hb.hex() if h not in acc_db: break else: return web.Response(text='server full', status=500) acc_db[h] = filename local_fname = f'{conf.data_path}/{acc}/{h}_{filename}' ext = os.path.splitext(filename)[1] if conf.show_ext else '' os.fdopen(os.open(local_fname, os.O_WRONLY | os.O_CREAT, 0o600)).close() try: valid_file = await recv_file(file, local_fname) except IOError: return web.Response(text='internal io error', status=500) if valid_file: c = AES.new(conf.del_crypt_key, AES.MODE_CBC) hb = Padding.pad(hb, AES.block_size) del_h = (c.encrypt(hb) + c.iv).hex() return web.Response(text=f'{conf.protocol}://{acc}{conf.prefix}/{h[:conf.url_hash_len]}{ext}\n' f'{conf.protocol}://{acc}{conf.prefix}/del/{del_h}') os.unlink(local_fname) del acc_db[h] return web.Response(text=f'file is bigger than {sizeof_fmt(conf.max_filesize)}', status=413) async def recv_file(file, local_fname): size = 0 with io.BufferedWriter(open(local_fname, 'wb')) as f: while True: chunk = await file.read_chunk() if not chunk: return True size += len(chunk) if size > conf.max_filesize: return False f.write(chunk) async def handle_delete(req, acc, acc_db): chashiv = req.match_info.get('hash', 'x') if not set(chashiv).issubset(valid_hash_chars) or len(chashiv) != 64: return web.Response(text='invalid delete link', status=400) chashiv = bytes.fromhex(chashiv) c = AES.new(conf.del_crypt_key, AES.MODE_CBC, iv=chashiv[AES.block_size:]) fhash = c.decrypt(chashiv[:AES.block_size]) try: fhash = Padding.unpad(fhash, AES.block_size).hex() except ValueError: pass if fhash not in acc_db or len(fhash) == 32: return web.Response(text='this file doesn\'t exist on the server', status=404) os.unlink(f"{conf.data_path}/{acc}/{fhash}_{acc_db[fhash]}") del acc_db[fhash] return web.Response(text='file deleted') async def handle_download(req, acc, acc_db): fhash = req.match_info.get('hash', '').split('.', 1)[0] if fhash not in acc_db: return web.Response(text='file not found', status=404) return web.FileResponse(f"{conf.data_path}/{acc}/{fhash}_{acc_db[fhash]}", headers={ 'CONTENT-DISPOSITION': f'inline;filename="{acc_db[fhash]}"' }) def main(): if not os.path.isdir(conf.data_path): os.mkdir(conf.data_path) for acc in os.listdir(conf.data_path): if not os.path.isdir(f'{conf.data_path}/{acc}'): continue for file in os.listdir(f"{conf.data_path}/{acc}"): if "_" in file: fhash, fname = file.split('_', 1) file_db[acc][fhash] = fname app = web.Application(middlewares=[prepare]) app.router.add_post(conf.prefix + '/post/{acc}', handle_upload) app.router.add_get(conf.prefix + '/del/{hash}/{acc}', handle_delete) app.router.add_get(conf.prefix + '/{hash}/{acc}', handle_download) web.run_app(app, port=80) if __name__ == '__main__': valid_hash_chars = set(string.hexdigits) file_db = defaultdict(dict) confname = sys.argv[1] if sys.argv[1:] and os.path.isfile(sys.argv[1]) else 'config.yaml' with open(confname) as cf: conf = cast_to_ad(yaml.load(cf)) if conf.url_hash_len > 31: raise ValueError('url_hash_len can\'t be bigger than 31') if not set(conf.max_filesize.replace(' ', ''))\ .issubset(valid_hash_chars | {'*'}): raise ValueError('max_filsize only can contain numbers and *') conf.max_filesize = eval(conf.max_filesize) conf.auth_tokens = set(conf.tokens) conf.prefix = conf.prefix.strip("/") if conf.prefix: conf.prefix = f'/{conf.prefix}' conf.del_crypt_key = hashlib.md5(conf.del_crypt_key.encode()).digest()[:16] main()