94 lines
2.7 KiB
Python
94 lines
2.7 KiB
Python
"""
|
|
Prosody mod_http_upload_external -> S3 handler.
|
|
|
|
Implements the same HMAC protocol as prosody-filer but stores files in S3.
|
|
|
|
Flow:
|
|
1. XMPP client asks Prosody for an upload slot
|
|
2. Prosody generates a signed PUT URL pointing here
|
|
3. Client PUTs the file; this handler verifies the HMAC and uploads to S3
|
|
4. On GET, handler generates a presigned S3 URL and redirects
|
|
"""
|
|
|
|
import hashlib
|
|
import hmac as hmac_mod
|
|
import logging
|
|
import os
|
|
|
|
import boto3
|
|
from flask import Flask, Response, redirect, request
|
|
|
|
app = Flask(__name__)
|
|
logging.basicConfig(level=logging.INFO)
|
|
log = logging.getLogger("s3-upload-handler")
|
|
|
|
SECRET = os.environ["UPLOAD_SECRET"]
|
|
S3_BUCKET = os.environ["S3_BUCKET"]
|
|
S3_REGION = os.environ.get("S3_REGION", "us-east-1")
|
|
S3_ENDPOINT = os.environ.get("S3_ENDPOINT")
|
|
PRESIGN_EXPIRE = int(os.environ.get("PRESIGN_EXPIRE", "3600"))
|
|
|
|
s3_kwargs = {
|
|
"region_name": S3_REGION,
|
|
}
|
|
if S3_ENDPOINT:
|
|
s3_kwargs["endpoint_url"] = S3_ENDPOINT
|
|
|
|
s3 = boto3.client("s3", **s3_kwargs)
|
|
|
|
|
|
def verify_token(path: str, size: str, content_type: str, token: str) -> bool:
|
|
"""Verify HMAC-SHA256 token from Prosody mod_http_upload_external."""
|
|
message = f"{path} {size} {content_type}"
|
|
expected = hmac_mod.new(
|
|
SECRET.encode(), message.encode(), hashlib.sha256
|
|
).hexdigest()
|
|
return hmac_mod.compare_digest(expected, token)
|
|
|
|
|
|
@app.route("/upload/<path:file_path>", methods=["PUT"])
|
|
def upload(file_path):
|
|
token = request.args.get("v", "")
|
|
content_type = request.args.get("t", "application/octet-stream")
|
|
file_size = request.args.get("s", "0")
|
|
|
|
hmac_path = f"upload/{file_path}"
|
|
|
|
if not verify_token(hmac_path, file_size, content_type, token):
|
|
log.warning("HMAC verification failed for %s", hmac_path)
|
|
return Response("Forbidden", status=403)
|
|
|
|
body = request.get_data()
|
|
|
|
if len(body) != int(file_size):
|
|
return Response("Content length mismatch", status=400)
|
|
|
|
s3.put_object(
|
|
Bucket=S3_BUCKET,
|
|
Key=file_path,
|
|
Body=body,
|
|
ContentType=content_type,
|
|
)
|
|
|
|
log.info("Uploaded %s (%s bytes) to s3://%s/%s", file_path, file_size, S3_BUCKET, file_path)
|
|
return Response(status=201)
|
|
|
|
|
|
@app.route("/upload/<path:file_path>", methods=["GET", "HEAD"])
|
|
def download(file_path):
|
|
try:
|
|
url = s3.generate_presigned_url(
|
|
"get_object",
|
|
Params={"Bucket": S3_BUCKET, "Key": file_path},
|
|
ExpiresIn=PRESIGN_EXPIRE,
|
|
)
|
|
except Exception:
|
|
log.exception("Failed to generate presigned URL for %s", file_path)
|
|
return Response("Not found", status=404)
|
|
|
|
return redirect(url, code=302)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app.run(host="0.0.0.0", port=5050)
|