pieter/app.py

import os
from flask import Flask, request
from pathlib import Path
from werkzeug.utils import secure_filename
import subprocess

UPLOAD_FOLDER = Path(os.environ['REPO_DIR']) / "pkgs"
API_KEY = os.environ['API_KEY']
ALLOWED_EXTENSIONS = {'pkg.tar.zst', 'pkg.tar.gz'}

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER


def require_api_key(func):
    def inner(*args, **kwargs):
        print(request.headers)
        if request.headers.get('x-api-key') and request.headers.get('x-api-key') == API_KEY:
            return func(*args, **kwargs)

        return {'message': 'Unauthorized.'}, 401

    return inner


def allowed_file(filename):
    return any(filename.endswith(ext) for ext in ALLOWED_EXTENSIONS)


@app.route('/publish', methods=['POST'])
@require_api_key
def upload_file():
    # check if the post request has the file part
    if 'file' not in request.files or not request.files['file']:
        return {'message': 'No file.'}, 400

    file = request.files['file']

    # Check wether a filename was provided
    if file.filename == '':
        return {'message': 'No filename.'}, 400

    # Check wether the filename has an allowed file extension
    if not allowed_file(file.filename):
        return {'message': 'Invalid extension.'}, 400

    # Create path for file & check if it already exists
    filename = secure_filename(file.filename)
    path = Path(app.config['UPLOAD_FOLDER']) / filename

    if path.exists():
        return {'message': 'File already exists.'}, 400

    path.parent.mkdir(parents=True, exist_ok=True)

    # Store the file
    file.save(path)

    # Run repo-add on the file
    res = subprocess.run(["repo-add", path.parent.parent / "pieter.db.tar.gz", path])

    if res.returncode != 0:
        path.unlink()

        return {'message': 'Failed to add file to repository.'}, 500

    print(f"Uploaded {file}")

    return {'message': 'Success.'}, 200