ShareX_Storage/sharex_server.py

173 lines
5.8 KiB
Python
Raw Normal View History

2018-06-10 15:49:53 +02:00
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
2018-06-10 15:49:53 +02:00
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())
2019-06-04 18:38:50 +02:00
@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
2018-06-10 15:49:53 +02:00
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)
2018-06-10 15:49:53 +02:00
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)
2018-06-10 15:49:53 +02:00
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]}"'
2018-06-10 15:49:53 +02:00
})
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
2018-06-10 15:49:53 +02:00
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'
2018-06-10 15:49:53 +02:00
with open(confname) as cf:
2019-06-04 18:38:50 +02:00
conf = AttrDict.from_dict_recur(yaml.load(cf))
2018-06-10 15:49:53 +02:00
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()