Spaces:
Sleeping
Sleeping
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,235 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import io
|
3 |
+
import json
|
4 |
+
import base64
|
5 |
+
import struct
|
6 |
+
import logging
|
7 |
+
|
8 |
+
# --- Python Web & Image Libraries ---
|
9 |
+
import gradio as gr
|
10 |
+
from fastapi import FastAPI, UploadFile, File, HTTPException
|
11 |
+
from pydantic import BaseModel
|
12 |
+
from PIL import Image
|
13 |
+
import numpy as np
|
14 |
+
|
15 |
+
# --- Cryptography Libraries ---
|
16 |
+
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
17 |
+
from cryptography.hazmat.primitives import hashes
|
18 |
+
from cryptography.hazmat.primitives import serialization
|
19 |
+
from cryptography.hazmat.primitives.asymmetric import padding
|
20 |
+
from cryptography.exceptions import InvalidTag
|
21 |
+
|
22 |
+
# --- Configure Logging ---
|
23 |
+
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
24 |
+
logger = logging.getLogger(__name__)
|
25 |
+
|
26 |
+
# --- Load the Secret Private Key from Hugging Face Secrets ---
|
27 |
+
# This is the most important part for security.
|
28 |
+
PLUGIN_PRIVATE_KEY = os.environ.get('PLUGIN_PRIVATE_KEY')
|
29 |
+
|
30 |
+
# --- Cryptographic Constants (must match the encryption side) ---
|
31 |
+
AES_KEY_SIZE_CRYPTO = 32
|
32 |
+
AES_GCM_NONCE_SIZE_CRYPTO = 12
|
33 |
+
AES_GCM_TAG_SIZE_CRYPTO = 16
|
34 |
+
ENCRYPTED_AES_KEY_LEN_SIZE_CRYPTO = 4
|
35 |
+
|
36 |
+
# ==============================================================================
|
37 |
+
# DECRYPTION LOGIC (Python port of the core functionality)
|
38 |
+
# ==============================================================================
|
39 |
+
|
40 |
+
def load_rsa_private_key(pem_bytes: bytes) -> any:
|
41 |
+
"""Loads a PEM-formatted private key."""
|
42 |
+
try:
|
43 |
+
return serialization.load_pem_private_key(pem_bytes, password=None)
|
44 |
+
except (ValueError, TypeError) as e:
|
45 |
+
logger.error(f"Failed to load private key: {e}")
|
46 |
+
raise ValueError("The provided private key is invalid or malformed.")
|
47 |
+
|
48 |
+
def decrypt_hybrid_payload(crypto_payload: bytes, private_key_pem: str) -> bytes:
|
49 |
+
"""Decrypts the full RSA-AES hybrid payload."""
|
50 |
+
private_key = load_rsa_private_key(private_key_pem.encode('utf-8'))
|
51 |
+
|
52 |
+
rsa_key_size_bytes = private_key.key_size // 8
|
53 |
+
|
54 |
+
# 1. Parse the crypto payload
|
55 |
+
offset = ENCRYPTED_AES_KEY_LEN_SIZE_CRYPTO
|
56 |
+
encrypted_aes_key_len = struct.unpack('>I', crypto_payload[:offset])[0]
|
57 |
+
|
58 |
+
if encrypted_aes_key_len != rsa_key_size_bytes:
|
59 |
+
raise ValueError("Encrypted key length mismatch. Key pair is likely incorrect.")
|
60 |
+
|
61 |
+
encrypted_aes_key = crypto_payload[offset : offset + encrypted_aes_key_len]
|
62 |
+
offset += encrypted_aes_key_len
|
63 |
+
|
64 |
+
nonce = crypto_payload[offset : offset + AES_GCM_NONCE_SIZE_CRYPTO]
|
65 |
+
offset += AES_GCM_NONCE_SIZE_CRYPTO
|
66 |
+
|
67 |
+
ciphertext_with_tag = crypto_payload[offset:]
|
68 |
+
|
69 |
+
# 2. Decrypt the AES key with RSA
|
70 |
+
recovered_aes_key = private_key.decrypt(
|
71 |
+
encrypted_aes_key,
|
72 |
+
padding.OAEP(
|
73 |
+
mgf=padding.MGF1(algorithm=hashes.SHA256()),
|
74 |
+
algorithm=hashes.SHA256(),
|
75 |
+
label=None
|
76 |
+
)
|
77 |
+
)
|
78 |
+
|
79 |
+
# 3. Decrypt the data with the recovered AES key
|
80 |
+
aesgcm = AESGCM(recovered_aes_key)
|
81 |
+
try:
|
82 |
+
decrypted_data = aesgcm.decrypt(nonce, ciphertext_with_tag, None)
|
83 |
+
return decrypted_data
|
84 |
+
except InvalidTag:
|
85 |
+
raise ValueError("Decryption failed (InvalidTag). The data is corrupt or was modified.")
|
86 |
+
|
87 |
+
def decode_from_image_buffer(image_buffer: bytes, private_key_pem: str) -> dict:
|
88 |
+
"""Main function to extract and decrypt data from an image buffer."""
|
89 |
+
try:
|
90 |
+
img = Image.open(io.BytesIO(image_buffer))
|
91 |
+
# Ensure image is 3-channel RGB to match the encryption process
|
92 |
+
img_rgb = img.convert("RGB")
|
93 |
+
pixel_data = np.array(img_rgb).ravel()
|
94 |
+
except Exception as e:
|
95 |
+
raise ValueError(f"Failed to process image: {e}")
|
96 |
+
|
97 |
+
# 1. Extract LSB header to get data length
|
98 |
+
HEADER_BITS = 32
|
99 |
+
if pixel_data.size < HEADER_BITS:
|
100 |
+
raise ValueError("Image is too small to contain a valid LSB header.")
|
101 |
+
|
102 |
+
header_binary_string = "".join(str(pixel & 1) for pixel in pixel_data[:HEADER_BITS])
|
103 |
+
data_length_bytes = int(header_binary_string, 2).to_bytes(4, byteorder='big')
|
104 |
+
data_length = struct.unpack('>I', data_length_bytes)[0]
|
105 |
+
|
106 |
+
if data_length == 0:
|
107 |
+
return {}
|
108 |
+
|
109 |
+
# 2. Extract LSB data payload
|
110 |
+
data_bits_count = data_length * 8
|
111 |
+
start_offset = HEADER_BITS
|
112 |
+
end_offset = start_offset + data_bits_count
|
113 |
+
|
114 |
+
if pixel_data.size < end_offset:
|
115 |
+
raise ValueError("Image data is corrupt or truncated.")
|
116 |
+
|
117 |
+
data_binary_string = "".join(str(pixel & 1) for pixel in pixel_data[start_offset:end_offset])
|
118 |
+
crypto_payload = int(data_binary_string, 2).to_bytes(data_length, byteorder='big')
|
119 |
+
|
120 |
+
# 3. Decrypt the payload
|
121 |
+
decrypted_bytes = decrypt_hybrid_payload(crypto_payload, private_key_pem)
|
122 |
+
|
123 |
+
return json.loads(decrypted_bytes.decode('utf-8'))
|
124 |
+
|
125 |
+
|
126 |
+
# ==============================================================================
|
127 |
+
# FASTAPI APP AND ENDPOINTS
|
128 |
+
# ==============================================================================
|
129 |
+
|
130 |
+
app = FastAPI()
|
131 |
+
|
132 |
+
# Pydantic model for the /mcp endpoint's request body
|
133 |
+
class MCPRequest(BaseModel):
|
134 |
+
image_base64: str
|
135 |
+
|
136 |
+
@app.post("/api/decode")
|
137 |
+
async def handle_decode_api(file: UploadFile = File(...)):
|
138 |
+
"""Handles standard file uploads for decryption."""
|
139 |
+
logger.info("Received request on /api/decode")
|
140 |
+
if not PLUGIN_PRIVATE_KEY:
|
141 |
+
raise HTTPException(status_code=500, detail="Server is not configured with a private key.")
|
142 |
+
|
143 |
+
image_buffer = await file.read()
|
144 |
+
try:
|
145 |
+
decoded_data = decode_from_image_buffer(image_buffer, PLUGIN_PRIVATE_KEY)
|
146 |
+
return {"success": True, "data": decoded_data}
|
147 |
+
except Exception as e:
|
148 |
+
logger.error(f"API Decode Error: {e}", exc_info=True)
|
149 |
+
raise HTTPException(status_code=400, detail=str(e))
|
150 |
+
|
151 |
+
@app.post("/mcp/decode")
|
152 |
+
async def handle_decode_mcp(request: MCPRequest):
|
153 |
+
"""Handles base64-encoded image strings for programmatic use."""
|
154 |
+
logger.info("Received request on /mcp/decode")
|
155 |
+
if not PLUGIN_PRIVATE_KEY:
|
156 |
+
raise HTTPException(status_code=500, detail="Server is not configured with a private key.")
|
157 |
+
|
158 |
+
try:
|
159 |
+
image_buffer = base64.b64decode(request.image_base64)
|
160 |
+
decoded_data = decode_from_image_buffer(image_buffer, PLUGIN_PRIVATE_KEY)
|
161 |
+
return {"success": True, "data": decoded_data}
|
162 |
+
except Exception as e:
|
163 |
+
logger.error(f"MCP Decode Error: {e}", exc_info=True)
|
164 |
+
raise HTTPException(status_code=400, detail=str(e))
|
165 |
+
|
166 |
+
# ==============================================================================
|
167 |
+
# GRADIO INTERFACE (The Web UI)
|
168 |
+
# ==============================================================================
|
169 |
+
|
170 |
+
with gr.Blocks(title="Secure Decoder API") as demo:
|
171 |
+
gr.Markdown("# Secure KeyLock Decoder API (Python/Gradio)")
|
172 |
+
gr.Markdown(
|
173 |
+
"This is a server-side API that decrypts images created with the KeyLock service. "
|
174 |
+
"It holds a private key securely and exposes endpoints for decryption.\n\n"
|
175 |
+
"See the **API Documentation** tab for usage details."
|
176 |
+
)
|
177 |
+
|
178 |
+
with gr.Tabs():
|
179 |
+
with gr.TabItem("API Status"):
|
180 |
+
gr.Markdown("## Server Health Check")
|
181 |
+
key_status = "✅ Loaded successfully from secrets." if PLUGIN_PRIVATE_KEY else "❌ NOT FOUND. The API will not work. Set the `PLUGIN_PRIVATE_KEY` secret in the Space settings."
|
182 |
+
gr.Textbox(label="Private Key Status", value=key_status, interactive=False)
|
183 |
+
gr.Markdown("Visit `/docs` to see the auto-generated FastAPI documentation for the endpoints.")
|
184 |
+
|
185 |
+
with gr.TabItem("API Documentation"):
|
186 |
+
gr.Markdown(
|
187 |
+
"""
|
188 |
+
## Public Key
|
189 |
+
Use the following public key in the [KeyLock RSA-BU](https://huggingface.co/spaces/broadfield-dev/KeyLock-RSA-BU) app to encrypt data for this service.
|
190 |
+
```pem
|
191 |
+
-----BEGIN PUBLIC KEY-----
|
192 |
+
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAy2sWjB1iQ3vK03U7e/9E
|
193 |
+
O6J1K/s0tBq4Pz8F3r9/i8s7t9R1p8Z4Y6h4f4O7w9p9Z0c8t7m4J1e9g7K9m6f3
|
194 |
+
R1k3y7v1w0l7z6s5v2l8l4t9v8z7y6t5k2x1c9v7z3k1y9w8r5t3s1v9a8d7g6f5
|
195 |
+
e4d3c2b1a9f8e7d6c5b4a3f2e1d0c9b8a7f6e5d4c3b2a1f0e9d8c7b6a5f4e3d2
|
196 |
+
c1b0a9f8e7d6c5b4a3f2e1d0c9b8a7f6e5d4c3b2a1f0e9d8c7b6a5f4e3d2c1b0
|
197 |
+
a9f8e7d6c5b4a3f2e1d0c9b8a7f6e5d4c3b2a1f0e9d8c7b6a5f4e3d2c1b0a9f8
|
198 |
+
e7d6
|
199 |
+
-----END PUBLIC KEY-----
|
200 |
+
```
|
201 |
+
---
|
202 |
+
## Endpoint 1: `/api/decode`
|
203 |
+
Standard file upload endpoint.
|
204 |
+
- **Method**: `POST`
|
205 |
+
- **Body**: `multipart/form-data`
|
206 |
+
- **Field**: `file` (the image file)
|
207 |
+
|
208 |
+
**cURL Example:**
|
209 |
+
```bash
|
210 |
+
curl -X POST -F "file=@/path/to/your/encrypted_image.png" https://your-space-name.hf.space/api/decode
|
211 |
+
```
|
212 |
+
---
|
213 |
+
## Endpoint 2: `/mcp/decode`
|
214 |
+
Programmatic endpoint accepting a base64-encoded image.
|
215 |
+
- **Method**: `POST`
|
216 |
+
- **Body**: `JSON`
|
217 |
+
|
218 |
+
**JSON Payload Schema:**
|
219 |
+
```json
|
220 |
+
{
|
221 |
+
"image_base64": "iVBORw0KGgoAAAANSUhEUgA..."
|
222 |
+
}
|
223 |
+
```
|
224 |
+
**cURL Example:**
|
225 |
+
```bash
|
226 |
+
# Note: You must properly escape the JSON payload for your shell.
|
227 |
+
curl -X POST -H "Content-Type: application/json" \\
|
228 |
+
-d '{"image_base64": "'$(base64 -w 0 /path/to/your/encrypted_image.png)'"}' \\
|
229 |
+
https://your-space-name.hf.space/mcp/decode
|
230 |
+
```
|
231 |
+
"""
|
232 |
+
)
|
233 |
+
|
234 |
+
# Mount the FastAPI app to the Gradio Blocks UI
|
235 |
+
app = gr.mount_app_onto_block(app, demo, path="/")
|