upload/main.py
2024-08-23 12:28:12 +02:00

99 lines
3.0 KiB
Python
Executable File

import os, uuid, time, json, sqlite3, threading
from flask import Flask, request, redirect, send_from_directory, abort, render_template_string
app = Flask(__name__)
with open('config.json') as f:
config = json.load(f)
UPLOAD_FOLDER = config['upload_folder']
MAX_FILE_SIZE = config['max_file_size']
EXPIRATION_TIME = config['expiration_time']
DATA_FILE = config['data_file']
URL_PREFIX = config['base_url']
NAME = config['name']
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
def init_db():
with sqlite3.connect(DATA_FILE) as conn:
conn.execute('CREATE TABLE IF NOT EXISTS files (key TEXT PRIMARY KEY, type TEXT NOT NULL, path TEXT, url TEXT, expiry REAL)')
init_db()
def load_data():
with sqlite3.connect(DATA_FILE) as conn:
return {row[0]: {'type': row[1], 'path': row[2], 'url': row[3], 'expiry': row[4]} for row in conn.execute('SELECT * FROM files')}
def save_data():
with sqlite3.connect(DATA_FILE) as conn:
conn.execute('DELETE FROM files')
conn.executemany('INSERT INTO files VALUES (?, ?, ?, ?, ?)',
[(k, v['type'], v.get('path'), v.get('url'), v['expiry']) for k, v in data.items()])
data = load_data()
def cleanup_files():
while True:
time.sleep(3600)
now = time.time()
expired_keys = [k for k, v in data.items() if v['expiry'] < now]
for k in expired_keys:
if data[k]['type'] == 'file':
os.remove(data[k]['path'])
data.pop(k)
save_data()
threading.Thread(target=cleanup_files, daemon=True).start()
@app.route('/', methods=['GET'])
def landing_page():
with open('landing.html') as f:
html = f.read().format(
name=NAME,
base_url=URL_PREFIX.rstrip('/'),
expiration_days=EXPIRATION_TIME // 86400,
max_file_size=MAX_FILE_SIZE // 1048576
)
return render_template_string(html)
@app.route('/', methods=['POST'])
def upload_file():
file = request.files.get('file')
if not file or file.content_length > MAX_FILE_SIZE:
return "No file uploaded or file too large\n", 400
key = uuid.uuid4().hex[:6]
filepath = os.path.join(UPLOAD_FOLDER, f"{key}_{file.filename}")
file.save(filepath)
data[key] = {'type': 'file', 'path': filepath, 'expiry': time.time() + EXPIRATION_TIME}
save_data()
return f"{URL_PREFIX}{key}\n"
@app.route('/shorten', methods=['POST'])
def shorten_url():
url = request.form.get('url')
if not url:
return "No URL provided\n", 400
key = uuid.uuid4().hex[:6]
data[key] = {'type': 'url', 'url': url, 'expiry': time.time() + EXPIRATION_TIME}
save_data()
return f"{URL_PREFIX}{key}\n"
@app.route('/<key>', methods=['GET'])
def get_content(key):
item = data.get(key)
if not item:
abort(404)
if item['type'] == 'file':
return send_from_directory(UPLOAD_FOLDER, os.path.basename(item['path']))
if item['type'] == 'url':
return redirect(item['url'])
abort(404)
if __name__ == '__main__':
app.run(host='0.0.0.0', port=5000)