broadfield-dev commited on
Commit
d123ded
·
verified ·
1 Parent(s): fd9604a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +235 -0
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="/")