Create core.py

core.py

@@ -0,0 +1,109 @@
+import io
+import json
+import os
+import struct
+import logging
+
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa, padding
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+
+# --- Configure Logging ---
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# --- Constants ---
+AES_KEY_SIZE_CRYPTO = 32
+AES_GCM_NONCE_SIZE_CRYPTO = 12
+LENGTH_HEADER_SIZE = 4
+
+def parse_kv_string_to_dict(kv_str: str) -> dict:
+    """Parses a string of key:value or key=value pairs into a dictionary."""
+    data_dict = {}
+    for line in kv_str.splitlines():
+        line = line.strip()
+        if not line or line.startswith('#'):
+            continue
+        if '=' in line:
+            key, value = line.split('=', 1)
+        elif ':' in line:
+            key, value = line.split(':', 1)
+        else:
+            raise ValueError(f"Invalid format on line: '{line}'. Use 'key: value' or 'key = value'.")
+        data_dict[key.strip()] = value.strip().strip("'\"")
+    return data_dict
+
+def encrypt_data_hybrid(data: bytes, public_key_pem: str) -> bytes:
+    """Encrypts data using RSA-OAEP + AES-GCM hybrid encryption."""
+    public_key = serialization.load_pem_public_key(public_key_pem.encode('utf-8'))
+    aes_key = os.urandom(AES_KEY_SIZE_CRYPTO)
+    nonce = os.urandom(AES_GCM_NONCE_SIZE_CRYPTO)
+    
+    ciphertext_with_tag = AESGCM(aes_key).encrypt(nonce, data, None)
+    
+    rsa_encrypted_aes_key = public_key.encrypt(
+        aes_key,
+        padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
+    )
+    
+    encrypted_aes_key_len_bytes = struct.pack('>I', len(rsa_encrypted_aes_key))
+    return encrypted_aes_key_len_bytes + rsa_encrypted_aes_key + nonce + ciphertext_with_tag
+
+def embed_data_in_image(img_obj: Image.Image, data: bytes) -> Image.Image:
+    """Embeds data into the LSB of an image."""
+    img_rgb = img_obj.convert("RGB")
+    pixel_data = np.array(img_rgb).ravel()
+    
+    data_length_header = struct.pack('>I', len(data))
+    binary_payload = ''.join(format(byte, '08b') for byte in data_length_header + data)
+    
+    if len(binary_payload) > pixel_data.size:
+        raise ValueError(f"Data too large for image capacity. Needs {len(binary_payload)} bits, image has {pixel_data.size}.")
+        
+    for i in range(len(binary_payload)):
+        pixel_data[i] = (pixel_data[i] & 0xFE) | int(binary_payload[i])
+        
+    stego_pixels = pixel_data.reshape(img_rgb.size[1], img_rgb.size[0], 3)
+    return Image.fromarray(stego_pixels, 'RGB')
+
+def create_carrier_image(width=800, height=600) -> Image.Image:
+    """Generates a simple carrier image."""
+    img = Image.new('RGB', (width, height), color = (73, 109, 137))
+    d = ImageDraw.Draw(img)
+    try:
+        font = ImageFont.truetype("DejaVuSans.ttf", 40)
+    except IOError:
+        font = ImageFont.load_default()
+    d.text((width/2, height/2), "KeyLock Encrypted Data", fill=(255,255,0), font=font, anchor="ms")
+    return img
+
+def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.Image:
+    """
+    High-level orchestrator function to create the final encrypted image.
+    """
+    logger.info("Starting image creation process...")
+    
+    # 1. Parse and validate secret data
+    if not secret_data_str.strip():
+        raise ValueError("Secret data cannot be empty.")
+    secret_dict = parse_kv_string_to_dict(secret_data_str)
+    if not secret_dict:
+        raise ValueError("No valid key-value pairs found in secret data.")
+    
+    # 2. Serialize data to JSON bytes
+    json_bytes = json.dumps(secret_dict).encode('utf-8')
+    logger.info(f"Serialized {len(json_bytes)} bytes of data.")
+
+    # 3. Encrypt the data
+    encrypted_payload = encrypt_data_hybrid(json_bytes, public_key_pem)
+    logger.info(f"Encrypted payload created with size: {len(encrypted_payload)} bytes.")
+
+    # 4. Create a carrier image and embed the data
+    carrier_image = create_carrier_image()
+    stego_image = embed_data_in_image(carrier_image, encrypted_payload)
+    logger.info("Data successfully embedded into carrier image.")
+    
+    return stego_image