Update client/app.py
Browse files- client/app.py +124 -274
client/app.py
CHANGED
@@ -1,295 +1,145 @@
|
|
1 |
-
import
|
2 |
-
from PIL import Image, ImageDraw, ImageFont, ImageOps
|
3 |
-
import base64
|
4 |
import io
|
5 |
import json
|
6 |
-
import
|
7 |
-
import os
|
8 |
-
import requests
|
9 |
import struct
|
10 |
-
import
|
|
|
|
|
|
|
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
|
17 |
-
from
|
|
|
|
|
|
|
|
|
18 |
|
19 |
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
20 |
logger = logging.getLogger(__name__)
|
21 |
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
T2I_PROMPT = "A stunning view of a distant galaxy, nebulae, and constellations, digital art, vibrant colors, cinematic lighting, 8k, masterpiece."
|
27 |
|
28 |
-
|
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 |
-
|
34 |
-
|
35 |
-
logger.
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
|
|
|
|
43 |
try:
|
44 |
-
|
45 |
-
|
46 |
-
|
47 |
-
|
48 |
-
|
49 |
-
|
|
|
50 |
except Exception as e:
|
51 |
-
logger.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
52 |
try:
|
53 |
-
|
54 |
-
|
55 |
-
|
56 |
-
|
57 |
-
|
58 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
59 |
except Exception as e:
|
60 |
-
logger.error(f"
|
61 |
-
raise gr.Error(f"
|
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
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|