Update app.py

app.py

@@ -1,63 +1,73 @@
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, send_file
 from flask_cors import CORS
-from datetime import datetime, timedelta
+from werkzeug.utils import secure_filename
+import tempfile
 import uuid
-import threading
+import os
+import io
+import base64
+from PIL import Image
 import time
 
 app = Flask(__name__)
 CORS(app)
 
-# In-memory storage
-storage = {}
+# In-memory store: { id: { data, image (optional), expire_at, view_once } }
+SECRETS = {}
+MAX_IMAGE_SIZE = 300 * 1024  # 300 KB
 
-@app.route("/")
-def index():
-    return "✅ Sharelock-style secret backend is running."
-
-@app.route('/api/store', methods=['POST'])
+@app.route("/api/store", methods=["POST"])
 def store():
-    req = request.get_json()
-    data = req.get('data')
-    ttl = int(req.get('ttl', 3600))  # Time to live in seconds (default: 1 hour)
-    view_once = bool(req.get('view_once', False))
-
-    secret_id = str(uuid.uuid4())[:8]
-    expire_at = datetime.utcnow() + timedelta(seconds=ttl)
-
-    storage[secret_id] = {
-        'data': data,
-        'expire_at': expire_at,
-        'view_once': view_once,
-    }
+    form = request.form
+    data = form.get("data")
+    ttl = int(form.get("ttl", 300))
+    view_once = form.get("view_once") == "true"
 
-    return jsonify({'id': secret_id})
+    # Handle image if present
+    image_file = request.files.get("image")
+    image_data = None
+
+    if image_file:
+        img_bytes = image_file.read()
+        if len(img_bytes) > MAX_IMAGE_SIZE:
+            image = Image.open(io.BytesIO(img_bytes))
+            image.thumbnail((1024, 1024))  # Resize for safety
+            output = io.BytesIO()
+            image.save(output, format="JPEG", optimize=True, quality=70)
+            image_data = base64.b64encode(output.getvalue()).decode("utf-8")
+        else:
+            image_data = base64.b64encode(img_bytes).decode("utf-8")
+
+    sid = str(uuid.uuid4())
+    SECRETS[sid] = {
+        "data": data,
+        "image": image_data,
+        "expire_at": time.time() + ttl,
+        "view_once": view_once
+    }
+    return jsonify({"id": sid})
 
-@app.route('/api/fetch/<secret_id>', methods=['GET'])
-def fetch(secret_id):
-    secret = storage.get(secret_id)
+@app.route("/api/fetch/<sid>")
+def fetch(sid):
+    secret = SECRETS.get(sid)
     if not secret:
-        return jsonify({'error': 'Not found'}), 404
-
-    if datetime.utcnow() > secret['expire_at']:
-        del storage[secret_id]
-        return jsonify({'error': 'Expired'}), 410
-
-    data = secret['data']
-    if secret['view_once']:
-        del storage[secret_id]
-    return jsonify({'data': data})
-
-def cleanup():
-    while True:
-        now = datetime.utcnow()
-        expired_keys = [k for k, v in storage.items() if now > v['expire_at']]
-        for k in expired_keys:
-            del storage[k]
-        time.sleep(60)  # Run every minute
-
-# Start cleanup thread
-threading.Thread(target=cleanup, daemon=True).start()
+        return jsonify({"error": "Not found"}), 404
+    if time.time() > secret["expire_at"]:
+        del SECRETS[sid]
+        return jsonify({"error": "Expired"}), 410
+
+    response = {"data": secret["data"]}
+    if secret.get("image"):
+        response["image"] = secret["image"]
+
+    if secret["view_once"]:
+        del SECRETS[sid]
+
+    return jsonify(response)
+
+@app.route("/")
+def index():
+    return "Sharelock Flask backend is running."
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)