#!/usr/bin/env python3

import os
import threading
import logging
import time
import traceback
import mimetypes
from hashlib import sha256
from pathlib import Path
from functools import partial

from flask import Flask, make_response, request, jsonify, abort
from sqlalchemy import (MetaData, Table, Column, CHAR, VARCHAR, DateTime,
                        LargeBinary, BOOLEAN)
from pyaltt2.config import load_yaml, config_value
from pyaltt2.db import Database
from pyaltt2.res import ResourceStorage
from pyaltt2.converters import val_to_boolean
from datetime import datetime, timedelta
import pyaltt2.crypto as cr
import pytz

logger = logging.getLogger('gunicorn.error')

dir_me = Path(__file__).absolute().parents[1]

rs = ResourceStorage(mod='secureshare')
rq = partial(rs.get, resource_subdir='sql', ext='sql')

app = Flask(__name__)

config = load_yaml(os.getenv('SECURESHARE_CONFIG'))['secureshare']

UPLOAD_KEY = config['upload-key']

db = Database(config['db'], rq_func=rq)

EXTERNAL_URL = config.get('url', '')


def ok(data=None):
    result = {'ok': True}
    if data:
        result.update(data)
    return jsonify(result)


@app.route('/', methods=['GET'])
def index():
    return ';)'


@app.route('/ping', methods=['GET'])
def ping():
    return make_response('', 204)


@app.route('/u', methods=['POST'])
def upload():
    if request.headers.get('x-auth-key') != UPLOAD_KEY:
        return make_response('Invalid upload key', 403)
    f = request.files.get('file', request.form.get('file'))
    if f is None:
        return make_response('File not uploaded', 403)
    filename = request.form.get('fname', f.filename)
    if filename is None:
        return make_response('File name not specified', 403)
    esecs = int(request.form.get('expires', config['default-expires']))
    expires = (datetime.now() + timedelta(seconds=esecs)).replace(
        tzinfo=pytz.timezone(time.tzname[0]))
    data = f.stream.read()
    sha256sum_gen = sha256()
    sha256sum_gen.update(data)
    sha256sum = sha256sum_gen.hexdigest()
    received_sha256sum = request.form.get('sha256sum')
    if received_sha256sum and received_sha256sum != sha256sum:
        return make_response('Checksum does not match', 422)
    file_id = cr.gen_random_str(16)
    file_key = cr.gen_random_str(16)
    filename = os.path.basename(filename)
    oneshot = val_to_boolean(request.form.get('oneshot', False))
    engine = cr.Rioja(file_key, bits=256)
    contents = engine.encrypt(data, b64=False)
    location = (f'{EXTERNAL_URL}/d/{file_id}/' f'{file_key}/{filename}')
    response = make_response(f'{location}\r\n', 201)
    response.headers['Location'] = location
    if EXTERNAL_URL:
        response.autocorrect_location_header = False
    response.headers['Cache-Control'] = ('no-cache, no-store, must-revalidate,'
                                         ' post-check=0, pre-check=0')
    response.headers['Pragma'] = 'no-cache'
    response.headers['Expires'] = expires.isoformat() + 'Z'
    mimetype = mimetypes.guess_type(filename)[0]
    if mimetype is None:
        mimetype = 'application/octet-stream'
    if mimetype == 'application/octet-stream':
        try:
            data.decode()
            mimetype = 'text/plain'
        except:
            pass
    db.query('stor.add',
             id=file_id,
             fname=filename,
             sha256sum=sha256sum,
             mimetype=mimetype,
             expires=expires,
             oneshot=oneshot,
             data=contents)
    return response


@app.route('/d/<file_id>/<file_key>/<file_name>', methods=['GET'])
def download(file_id, file_key, file_name):
    delete = request.args.get('c') == 'rm'
    try:
        f = db.qlookup('stor.get', id=file_id, fname=file_name)
    except LookupError:
        abort(404)
    engine = cr.Rioja(file_key, bits=256)
    try:
        contents = engine.decrypt(f['data'].tobytes(), b64=False)
    except ValueError:
        abort(403)
    if delete or f['oneshot']:
        db.query('stor.delete', id=file_id)
    if delete:
        return ok()
    response = make_response(contents)
    response.headers['Content-Type'] = f['mimetype']
    response.headers['x-hash-sha256'] = f['sha256sum']
    if val_to_boolean(request.args.get('raw')):
        response.headers[
            'Content-Disposition'] = f'attachment;filename={file_name}'
    return response


def clean_db():
    while True:
        logger.debug('cleaner worker running')
        try:
            db.query('delete.expired', d=datetime.now())
        except:
            logger.error(traceback.format_exc())
        time.sleep(config.get('db-clean-interval', 60))


dbconn = db.connect()

meta = MetaData()
stor = Table('stor', meta, Column('id', CHAR(16), primary_key=True),
             Column('fname', VARCHAR(255), nullable=False),
             Column('sha256sum', CHAR(64), nullable=False),
             Column('mimetype', VARCHAR(255), nullable=False),
             Column('expires', DateTime(timezone=True), nullable=False),
             Column('oneshot', BOOLEAN, nullable=False, server_default='0'),
             Column('data', LargeBinary, nullable=False))

meta.create_all(dbconn)

dbconn.close()

threading.Thread(target=clean_db, daemon=True).start()
