broadfield-dev commited on
Commit
1049d4a
·
verified ·
1 Parent(s): 5ba8863

Update client/app.py

Browse files
Files changed (1) hide show
  1. client/app.py +124 -274
client/app.py CHANGED
@@ -1,295 +1,145 @@
1
- import gradio as gr
2
- from PIL import Image, ImageDraw, ImageFont, ImageOps
3
- import base64
4
  import io
5
  import json
6
- import logging
7
- import os
8
- import requests
9
  import struct
10
- import tempfile # <--- Import tempfile
 
 
 
11
  import numpy as np
12
- from cryptography.hazmat.primitives import serialization
13
- from cryptography.hazmat.primitives.asymmetric import rsa, padding
14
  from cryptography.hazmat.primitives.ciphers.aead import AESGCM
15
  from cryptography.hazmat.primitives import hashes
16
- from gradio_client import Client, handle_file
17
- from huggingface_hub import InferenceClient
 
 
 
 
18
 
19
  logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
20
  logger = logging.getLogger(__name__)
21
 
22
- # --- Constants ---
23
- DEFAULT_IMAGE_URL = "https://images.unsplash.com/photo-1506318137071-a8e063b4bec0?q=80&w=1200&auto=format&fit=crop"
24
- IMAGE_SIZE = 600
25
- T2I_MODEL = "sd-community/sdxl-lightning"
26
- T2I_PROMPT = "A stunning view of a distant galaxy, nebulae, and constellations, digital art, vibrant colors, cinematic lighting, 8k, masterpiece."
27
 
28
- # --- Helper Functions for Image Handling (Unchanged) ---
29
-
30
- def resize_and_crop(img: Image.Image, size: int = IMAGE_SIZE) -> Image.Image:
31
- """Resizes an image to fit within a square of `size` and then center-crops it."""
32
  try:
33
- return ImageOps.fit(img, (size, size), Image.Resampling.LANCZOS)
34
- except Exception as e:
35
- logger.error(f"Failed to resize and crop image: {e}")
36
- return img.resize((size, size), Image.Resampling.LANCZOS)
37
-
38
- def prepare_base_image(uploaded_image: Image.Image | None, progress) -> Image.Image:
39
- """Provides a base image using a fallback logic, updating a Gradio progress object."""
40
- if uploaded_image:
41
- progress(0, desc=" Using uploaded image...")
42
- return resize_and_crop(uploaded_image)
 
 
43
  try:
44
- progress(0, desc="⏳ Fetching default background...")
45
- response = requests.get(DEFAULT_IMAGE_URL, timeout=15)
46
- response.raise_for_status()
47
- img = Image.open(io.BytesIO(response.content)).convert("RGB")
48
- progress(0, desc="✅ Using default background image.")
49
- return resize_and_crop(img)
 
50
  except Exception as e:
51
- logger.warning(f"Could not fetch default image: {e}. Falling back to AI generation.")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
52
  try:
53
- progress(0, desc=f"⏳ Generating image with {T2I_MODEL}...")
54
- client = InferenceClient()
55
- image_bytes = client.text_to_image(T2I_PROMPT, model=T2I_MODEL)
56
- img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
57
- progress(0, desc="✅ New image generated.")
58
- return resize_and_crop(img)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
59
  except Exception as e:
60
- logger.error(f"Fatal: All image sources failed. Text-to-image failed with: {e}")
61
- raise gr.Error(f"Failed to obtain a base image. AI generation error: {e}")
62
-
63
- # --- Core Cryptography and Image Creation (Unchanged) ---
64
 
65
  def generate_rsa_keys():
66
  private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
67
  private_pem = private_key.private_bytes(encoding=serialization.Encoding.PEM, format=serialization.PrivateFormat.PKCS8, encryption_algorithm=serialization.NoEncryption()).decode('utf-8')
68
  public_pem = private_key.public_key().public_bytes(encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo).decode('utf-8')
69
- return private_pem, public_pem
70
-
71
- def create_encrypted_image(secret_data_str: str, public_key_pem: str, base_image: Image.Image, overlay_option: str) -> Image.Image:
72
- if not secret_data_str.strip(): raise ValueError("Secret data cannot be empty.")
73
- if not public_key_pem.strip(): raise ValueError("Public Key cannot be empty.")
74
- data_dict = {}
75
- for line in secret_data_str.splitlines():
76
- line = line.strip()
77
- if not line or line.startswith('#'): continue
78
- parts = line.split(':', 1) if ':' in line else line.split('=', 1)
79
- if len(parts) == 2: data_dict[parts[0].strip()] = parts[1].strip().strip("'\"")
80
- if not data_dict: raise ValueError("No valid key-value pairs found.")
81
- json_bytes = json.dumps(data_dict).encode('utf-8')
82
- public_key = serialization.load_pem_public_key(public_key_pem.encode('utf-8'))
83
- aes_key, nonce = os.urandom(32), os.urandom(12)
84
- ciphertext = AESGCM(aes_key).encrypt(nonce, json_bytes, None)
85
- rsa_encrypted_key = public_key.encrypt(aes_key, padding.OAEP(mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None))
86
- encrypted_payload = struct.pack('>I', len(rsa_encrypted_key)) + rsa_encrypted_key + nonce + ciphertext
87
- img = base_image.copy().convert("RGB")
88
- width, height = img.size
89
- # (Overlay drawing code is unchanged and correct)
90
- draw = ImageDraw.Draw(img, "RGBA")
91
- try:
92
- font_bold = ImageFont.truetype("DejaVuSans-Bold.ttf", 30)
93
- font_regular = ImageFont.truetype("DejaVuSans.ttf", 15)
94
- except IOError:
95
- font_bold = ImageFont.load_default(size=28)
96
- font_regular = ImageFont.load_default(size=14)
97
- overlay_color, title_color = (15, 23, 42, 190), (226, 232, 240)
98
- key_color, value_color = (148, 163, 184), (241, 245, 249)
99
- draw.rectangle([0, 20, width, 80], fill=overlay_color)
100
- draw.text((width / 2, 50), "KeyLock Secure Data", fill=title_color, font=font_bold, anchor="ms")
101
- if overlay_option != "None":
102
- lines = list(data_dict.keys()) if overlay_option == "Keys Only" else [f"{k}: {v}" for k, v in data_dict.items()]
103
- line_heights = [draw.textbbox((0, 0), line, font=font_regular)[3] for line in lines]
104
- total_text_height = sum(line_heights) + (len(lines) - 1) * 6
105
- box_height = total_text_height + 30
106
- box_y0 = height - box_height - 20
107
- draw.rectangle([20, box_y0, width - 20, height - 20], fill=overlay_color)
108
- current_y = box_y0 + 15
109
- for i, (key, value) in enumerate(data_dict.items()):
110
- if overlay_option == "Keys Only": draw.text((35, current_y), key, fill=key_color, font=font_regular)
111
- else:
112
- key_text = f"{key}:"; draw.text((35, current_y), key_text, fill=key_color, font=font_regular)
113
- key_bbox = draw.textbbox((35, current_y), key_text, font=font_regular)
114
- draw.text((key_bbox[2] + 8, current_y), str(value), fill=value_color, font=font_regular)
115
- current_y += line_heights[i] + 6
116
- # (Steganography code is unchanged and correct)
117
- pixel_data = np.array(img.convert("RGB")).ravel()
118
- header = struct.pack('>I', len(encrypted_payload))
119
- binary_payload = ''.join(format(b, '08b') for b in header + encrypted_payload)
120
- if len(binary_payload) > pixel_data.size: raise ValueError(f"Data is too large. Max: {pixel_data.size // 8} bytes.")
121
- pixel_data[:len(binary_payload)] = (pixel_data[:len(binary_payload)] & 0xFE) | np.array(list(binary_payload), dtype=np.uint8)
122
- stego_pixels = pixel_data.reshape((height, width, 3))
123
- return Image.fromarray(stego_pixels, 'RGB')
124
-
125
- # --- New Server Management and Gradio Wrappers ---
126
-
127
- def add_server(url: str, current_servers: list):
128
- # This function is correct from the previous fix. Unchanged.
129
- if not url or not url.startswith(('http://', 'https://')): raise gr.Error("Please enter a valid server URL (e.g., https://...).")
130
- url = url.strip().rstrip('/')
131
- if any(s['url'] == url for s in current_servers):
132
- gr.Info(f"Server at {url} is already in the list.")
133
- return current_servers, f"ℹ️ Server at {url} is already in the list.", gr.Button(interactive=True)
134
- status = f"⏳ Testing server at {url}..."
135
- yield current_servers, status, gr.Button(interactive=False)
136
- try:
137
- client = Client(url, verbose=False)
138
- server_info = client.predict(api_name="/keylock-info")
139
- if not isinstance(server_info, dict): raise TypeError(f"Expected dict from /keylock-info, but got {type(server_info)}")
140
- server_name = server_info.get("name", url)
141
- public_key = client.predict(api_name="/keylock-pub")
142
- if not isinstance(public_key, str): raise TypeError(f"Expected a string from /keylock-pub, but got {type(public_key)}")
143
- if not public_key.strip().startswith("-----BEGIN PUBLIC KEY-----"): raise ValueError("Received invalid public key format from server.")
144
- new_server = {"name": f"{server_name} ({url.split('//')[1].split('/')[0]})", "url": url, "public_key": public_key, "api_endpoint": "/keylock-server"}
145
- updated_servers = current_servers + [new_server]
146
- status_message = f"✅ Successfully added server: {server_name}"; gr.Info(status_message)
147
- yield updated_servers, status_message, gr.Button(interactive=True)
148
- except Exception as e:
149
- logger.error(f"Failed to add server at {url}: {e}", exc_info=True)
150
- status_message = f"❌ Failed to connect or validate server at {url}. Error: {e}"; gr.Error(status_message)
151
- yield current_servers, status_message, gr.Button(interactive=True)
152
-
153
- def create_keylock_wrapper(service_name: str, secret_data: str, available_servers: list, uploaded_image: Image.Image | None, overlay_option: str, progress=gr.Progress(track_tqdm=True)):
154
- """
155
- Creates the encrypted image and returns it to two separate components:
156
- 1. A gr.Image for visual preview (which may be re-encoded by Gradio).
157
- 2. A gr.File for a guaranteed, uncorrupted PNG download.
158
- """
159
- if not service_name: raise gr.Error("Please add and select a target server.")
160
- server = next((s for s in available_servers if s['name'] == service_name), None)
161
- if not server: raise gr.Error(f"Could not find config for '{service_name}'. Please re-add it.")
162
-
163
- png_path = None # Initialize to handle potential errors
164
- try:
165
- base_image = prepare_base_image(uploaded_image, progress)
166
- progress(0.5, desc="Encrypting and embedding data...")
167
- created_image = create_encrypted_image(secret_data, server['public_key'], base_image, overlay_option)
168
-
169
- # Save image to a temporary PNG file to prevent re-encoding.
170
- # This temporary file will be served by gr.File for perfect data integrity.
171
- with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_f:
172
- png_path = temp_f.name
173
- created_image.save(png_path, "PNG", compress_level=1)
174
-
175
- status_message = f"✅ Success! Image created for '{service_name}'. Use the link to download the uncorrupted file."
176
- # Return path to both preview and file components. Make file component visible.
177
- return png_path, png_path, status_message, gr.Tabs(selected=1), gr.File(visible=True)
178
-
179
- except Exception as e:
180
- logger.error(f"Error creating image: {e}", exc_info=True)
181
- if png_path and os.path.exists(png_path):
182
- os.remove(png_path) # Clean up temp file on error
183
- # Return Nones and hide file component
184
- return None, None, f"❌ Error: {e}", gr.Tabs(), gr.File(visible=False)
185
-
186
- def send_keylock_wrapper(service_name: str, image_path: str, available_servers: list):
187
- # This function now expects a file path `image_path` from gr.Image
188
- if not service_name: raise gr.Error("Please select a target server.")
189
- if image_path is None: raise gr.Error("Please create or upload an image to send.")
190
-
191
- server = next((s for s in available_servers if s['name'] == service_name), None)
192
- if not server: raise gr.Error(f"Config Error for '{service_name}'. Please re-add it.")
193
-
194
- server_url, api_name = server['url'], server['api_endpoint']
195
- yield None, f"⏳ Connecting to remote server: {server_url}"
196
- try:
197
- # The input `image_path` is a path to a temporary file created by gr.Image.
198
- # We read it directly to ensure we get the correct bytes.
199
- with open(image_path, "rb") as img_file:
200
- b64_string = base64.b64encode(img_file.read()).decode("utf-8")
201
-
202
- client = Client(server_url)
203
- result = client.predict(image_base64_string=b64_string, api_name=api_name)
204
-
205
- if isinstance(result, str) and os.path.exists(result):
206
- with open(result, 'r', encoding='utf-8') as f: decrypted_data = json.load(f)
207
- else: decrypted_data = result
208
-
209
- yield decrypted_data, "✅ Success! Data decrypted by remote server."
210
- except Exception as e:
211
- logger.error(f"Error sending to server: {e}", exc_info=True)
212
- yield None, f"❌ Error calling server API: {e}"
213
-
214
- # --- Gradio UI ---
215
- theme = gr.themes.Soft(primary_hue="blue", secondary_hue="sky", neutral_hue="slate")
216
-
217
- with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
218
- servers_state = gr.State([])
219
-
220
- gr.Markdown("# 🔑 KeyLock Operations Dashboard")
221
- gr.Markdown("A client to discover KeyLock servers, create encrypted images, and test decryption against live APIs.")
222
-
223
- with gr.Accordion("🔗 Add a KeyLock Server", open=True):
224
- with gr.Row():
225
- server_url_input = gr.Textbox(label="Server Base URL", placeholder="https://your-hf-space-name.hf.space")
226
- add_server_button = gr.Button("Test & Add Server", variant="secondary")
227
- global_status = gr.Textbox(label="Status", interactive=False, lines=1, value="Enter a server URL to begin.")
228
-
229
- with gr.Tabs() as tabs:
230
- with gr.TabItem("① Create KeyLock", id=0):
231
- with gr.Row(variant="panel"):
232
- with gr.Column(scale=2):
233
- creator_service_dropdown = gr.Dropdown(label="Target Server", interactive=True, info="Select a server to encrypt data for.")
234
- creator_secret_input = gr.Textbox(lines=5, label="Secret Data (key:value format)", placeholder="API_KEY: sk-123...\nUSER: demo-user")
235
- creator_base_image_input = gr.Image(label="Optional Base Image (600x600 recommended)", type="pil", sources=["upload"])
236
- creator_overlay_option = gr.Radio(label="Overlay Content", choices=["Keys and Values", "Keys Only", "None"], value="Keys and Values")
237
- creator_button = gr.Button("✨ Create Auth Image", variant="primary")
238
- with gr.Column(scale=3):
239
- # Use two components: one for preview, one for lossless download
240
- creator_image_preview = gr.Image(label="Image Preview (for display only)", type="filepath", height=500)
241
- creator_file_output = gr.File(label="Download Uncorrupted PNG File", file_count="single", visible=False)
242
-
243
- with gr.TabItem("② Send KeyLock", id=1):
244
- with gr.Row(variant="panel"):
245
- with gr.Column(scale=1):
246
- send_service_dropdown = gr.Dropdown(label="Target Server", interactive=True)
247
- # This component will receive the file path from the preview image
248
- client_image_input = gr.Image(type="filepath", label="Upload or Drag Encrypted Image", sources=["upload", "clipboard"])
249
- client_button = gr.Button("🔓 Decrypt via Remote Server", variant="primary")
250
- with gr.Column(scale=1):
251
- client_json_output = gr.JSON(label="Decrypted Data from Server")
252
-
253
- with gr.TabItem("ℹ️ Key Generation", id=2):
254
- with gr.Accordion("🔑 RSA Key Pair Generator", open=True):
255
- gr.Markdown("Use this utility to generate a new key pair for a new server. The server administrator must configure the server with the private key.")
256
- gen_keys_button = gr.Button("⚙️ Generate New 2048-bit Key Pair")
257
- output_private_key = gr.Textbox(lines=10, label="Generated Private Key (KEEP THIS SECRET!)", interactive=False, show_copy_button=True)
258
- output_public_key = gr.Textbox(lines=10, label="Generated Public Key", interactive=False, show_copy_button=True)
259
-
260
- # --- Event Handlers ---
261
- def sync_dropdowns(x): return x
262
-
263
- def update_dropdowns_from_state(servers_list):
264
- server_names = [s['name'] for s in servers_list] if servers_list else []
265
- selected = server_names[-1] if server_names else None
266
- dropdown_update = gr.Dropdown(choices=server_names, value=selected, label="Target Server", interactive=True)
267
- return dropdown_update, dropdown_update
268
-
269
- add_server_button.click(fn=add_server, inputs=[server_url_input, servers_state], outputs=[servers_state, global_status, add_server_button]).then(fn=lambda: "", outputs=[server_url_input])
270
-
271
- servers_state.change(fn=update_dropdowns_from_state, inputs=servers_state, outputs=[creator_service_dropdown, send_service_dropdown])
272
-
273
- gen_keys_button.click(fn=generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
274
-
275
- creator_button.click(
276
- fn=create_keylock_wrapper,
277
- inputs=[creator_service_dropdown, creator_secret_input, servers_state, creator_base_image_input, creator_overlay_option],
278
- outputs=[creator_image_preview, creator_file_output, global_status, tabs, creator_file_output]
279
- )
280
-
281
- client_button.click(
282
- fn=send_keylock_wrapper,
283
- inputs=[send_service_dropdown, client_image_input, servers_state],
284
- outputs=[client_json_output, global_status]
285
- )
286
-
287
- creator_service_dropdown.change(fn=sync_dropdowns, inputs=creator_service_dropdown, outputs=send_service_dropdown)
288
- send_service_dropdown.change(fn=sync_dropdowns, inputs=send_service_dropdown, outputs=creator_service_dropdown)
289
-
290
- # When a new image preview is generated, auto-populate it in the sender tab.
291
- # This works because both components now use file paths.
292
- creator_image_preview.change(fn=lambda x: x, inputs=creator_image_preview, outputs=client_image_input)
293
-
294
- if __name__ == "__main__":
295
- demo.launch()
 
1
+ import os
 
 
2
  import io
3
  import json
4
+ import base64
 
 
5
  import struct
6
+ import logging
7
+
8
+ import gradio as gr
9
+ from PIL import Image
10
  import numpy as np
 
 
11
  from cryptography.hazmat.primitives.ciphers.aead import AESGCM
12
  from cryptography.hazmat.primitives import hashes
13
+ from cryptography.hazmat.primitives import serialization
14
+ from cryptography.hazmat.primitives.asymmetric import padding, rsa
15
+ from cryptography.exceptions import InvalidTag
16
+
17
+ HEADER_BITS = 32
18
+ AES_GCM_NONCE_SIZE = 12
19
 
20
  logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
21
  logger = logging.getLogger(__name__)
22
 
23
+ KEYLOCK_PRIV_KEY_PEM = os.environ.get('KEYLOCK_PRIV_KEY')
24
+ PRIVATE_KEY_OBJECT = None
25
+ PUBLIC_KEY_PEM_STRING = ""
26
+ KEYLOCK_STATUS_MESSAGE = ""
 
27
 
28
+ if not KEYLOCK_PRIV_KEY_PEM:
29
+ dev_key_path = os.path.join(os.path.dirname(__file__), '..', 'keys', 'DEMO_ONLY_THIS _IS_SECRET_keylock_priv_key.pem')
 
 
30
  try:
31
+ with open(dev_key_path, "r") as f:
32
+ KEYLOCK_PRIV_KEY_PEM = f.read()
33
+ logger.warning(f"Loaded private key from dev path. This is for local testing only.")
34
+ KEYLOCK_STATUS_MESSAGE = f"⚠️ Loaded from development key file. This is for local testing but insecure for production."
35
+ except FileNotFoundError:
36
+ logger.error(f"FATAL: Private key not found at '{dev_key_path}' and 'KEYLOCK_PRIV_KEY' secret is not set.")
37
+ KEYLOCK_STATUS_MESSAGE = " NOT FOUND. The API is non-functional. Set the `KEYLOCK_PRIV_KEY` secret or provide the demo key file."
38
+ else:
39
+ logger.info("Successfully loaded private key from environment variable 'KEYLOCK_PRIV_KEY'.")
40
+ KEYLOCK_STATUS_MESSAGE = "✅ Loaded successfully from secrets/environment variable. Recommended secure configuration."
41
+
42
+ if KEYLOCK_PRIV_KEY_PEM:
43
  try:
44
+ PRIVATE_KEY_OBJECT = serialization.load_pem_private_key(KEYLOCK_PRIV_KEY_PEM.encode(), password=None)
45
+ public_key = PRIVATE_KEY_OBJECT.public_key()
46
+ PUBLIC_KEY_PEM_STRING = public_key.public_bytes(
47
+ encoding=serialization.Encoding.PEM,
48
+ format=serialization.PublicFormat.SubjectPublicKeyInfo
49
+ ).decode('utf-8')
50
+ KEYLOCK_STATUS_MESSAGE += "\n✅ Public key derived successfully."
51
  except Exception as e:
52
+ logger.error(f"Failed to load private key or derive public key: {e}", exc_info=True)
53
+ PRIVATE_KEY_OBJECT = None
54
+ PUBLIC_KEY_PEM_STRING = f"Error: Could not process the configured private key. Details: {e}"
55
+ KEYLOCK_STATUS_MESSAGE += f"\n❌ Failed to parse key: {e}"
56
+
57
+ MOCK_USER_DATABASE = {
58
+ "sk-12345-abcde": {"user": "demo-user", "permissions": "read"},
59
+ "sk-67890-fghij": {"user": "admin-user", "permissions": "read,write,delete"}
60
+ }
61
+
62
+ def example_authenticate(api_key: str, user_id: str) -> bool:
63
+ if not api_key or not user_id:
64
+ return False
65
+ db_entry = MOCK_USER_DATABASE.get(api_key)
66
+ if db_entry and db_entry.get("user") == user_id:
67
+ logger.info(f"Authentication successful for user '{user_id}'.")
68
+ return True
69
+ else:
70
+ logger.warning(f"Authentication failed for user '{user_id}' with key '{api_key[:8]}...'.")
71
+ return False
72
+
73
+ def get_public_key():
74
+ if not PUBLIC_KEY_PEM_STRING or "Error" in PUBLIC_KEY_PEM_STRING:
75
+ raise gr.Error("Server key is not configured correctly.")
76
+ return PUBLIC_KEY_PEM_STRING
77
+
78
+ def get_server_info():
79
+ return {
80
+ "name": "KeyLock Auth Server",
81
+ "version": "1.3",
82
+ "documentation": "This server decrypts data hidden in KeyLock images and performs a mock authentication. Use /keylock-pub to get the public key, and POST to /keylock-server with a base64 image string to decrypt and authenticate.",
83
+ "endpoints": {
84
+ "/keylock-pub": "GET - Returns the server's public key.",
85
+ "/keylock-info": "GET - Returns this information object.",
86
+ "/keylock-server": "POST - Decrypts a KeyLock image and attempts authentication."
87
+ }
88
+ }
89
+
90
+ def decode_data(image_base64_string: str) -> dict:
91
+ if not PRIVATE_KEY_OBJECT:
92
+ error_msg = "Server Error: The API is not configured with a private key."
93
+ logger.error(error_msg)
94
+ raise gr.Error(error_msg)
95
  try:
96
+ image_buffer = base64.b64decode(image_base64_string)
97
+ img = Image.open(io.BytesIO(image_buffer)).convert("RGB")
98
+ pixel_data = np.array(img).ravel()
99
+ header_binary_string = "".join(str(pixel & 1) for pixel in pixel_data[:HEADER_BITS])
100
+ data_length = int(header_binary_string, 2)
101
+ if data_length == 0: raise ValueError("No data found in image.")
102
+ data_bits_count = data_length * 8
103
+ end_offset = HEADER_BITS + data_bits_count
104
+ if pixel_data.size < end_offset: raise ValueError("Image data corrupt or truncated.")
105
+ data_binary_string = "".join(str(pixel & 1) for pixel in pixel_data[HEADER_BITS:end_offset])
106
+ crypto_payload = int(data_binary_string, 2).to_bytes(data_length, byteorder='big')
107
+ offset = 4
108
+ encrypted_aes_key_len = struct.unpack('>I', crypto_payload[:offset])[0]
109
+ encrypted_aes_key = crypto_payload[offset : offset + encrypted_aes_key_len]
110
+ offset += encrypted_aes_key_len
111
+ nonce = crypto_payload[offset : offset + AES_GCM_NONCE_SIZE]
112
+ offset += AES_GCM_NONCE_SIZE
113
+ ciphertext_with_tag = crypto_payload[offset:]
114
+ recovered_aes_key = PRIVATE_KEY_OBJECT.decrypt(encrypted_aes_key, padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None))
115
+ decrypted_bytes = AESGCM(recovered_aes_key).decrypt(nonce, ciphertext_with_tag, None)
116
+ decrypted_payload = json.loads(decrypted_bytes.decode('utf-8'))
117
+ logger.info(f"Successfully decoded payload: {decrypted_payload}")
118
+ api_key = decrypted_payload.get('API_KEY')
119
+ user_id = decrypted_payload.get('USER')
120
+ is_authenticated = example_authenticate(api_key=api_key, user_id=user_id)
121
+ if is_authenticated:
122
+ return {
123
+ "authentication_status": "Success",
124
+ "message": f"User '{user_id}' successfully authenticated.",
125
+ "granted_permissions": MOCK_USER_DATABASE[api_key]['permissions'],
126
+ "decoded_payload": decrypted_payload
127
+ }
128
+ else:
129
+ return {
130
+ "authentication_status": "Failed",
131
+ "message": "Authentication failed. Invalid credentials provided in the image.",
132
+ "decoded_payload": decrypted_payload
133
+ }
134
+ except (ValueError, InvalidTag, TypeError, struct.error) as e:
135
+ logger.warning(f"Decryption failed: {e}")
136
+ raise gr.Error(f"Decryption failed. Image may be corrupt or used the wrong public key. Details: {e}")
137
  except Exception as e:
138
+ logger.error(f"An unexpected server error occurred during decryption: {e}", exc_info=True)
139
+ raise gr.Error(f"An unexpected server error occurred. Details: {e}")
 
 
140
 
141
  def generate_rsa_keys():
142
  private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
143
  private_pem = private_key.private_bytes(encoding=serialization.Encoding.PEM, format=serialization.PrivateFormat.PKCS8, encryption_algorithm=serialization.NoEncryption()).decode('utf-8')
144
  public_pem = private_key.public_key().public_bytes(encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo).decode('utf-8')
145
+ return private_pem, public_pem