ShareX_Storage/sharex_server.py

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, hdrs


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())

    @staticmethod
    def from_dict_recur(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] = AttrDict.from_dict_recur(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='url key-space 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={
        hdrs.CACHE_CONTROL: "no-cache",
        hdrs.CONTENT_DISPOSITION: f'inline;filename="{acc_db[fhash]}"'
    })


def main():
    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_filesize 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]
    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 = AttrDict.from_dict_recur(yaml.safe_load(cf))
    main()