image_kb_logic.py

import os
import io
import struct
import logging
import random
import json
from datetime import datetime
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.exceptions import InvalidTag

logger = logging.getLogger(__name__)

KEY_SIZE = 32
SALT_SIZE = 16
NONCE_SIZE = 12
TAG_SIZE = 16
PBKDF2_ITERATIONS = 480000
LENGTH_HEADER_SIZE = 4
PREFERRED_FONTS = ["Arial", "Helvetica", "DejaVu Sans", "Verdana", "Calibri", "sans-serif"]
MAX_KEYS_TO_DISPLAY_OVERLAY = 15

def convert_pil_to_png_bytes(image: Image.Image) -> bytes:
    with io.BytesIO() as buffer:
        image.save(buffer, format="PNG")
        return buffer.getvalue()

def _get_font(preferred_fonts, base_size):
    fp = None
    safe_base_size = int(base_size)
    if safe_base_size <= 0: safe_base_size = 10
    for n in preferred_fonts:
        try: ImageFont.truetype(n.lower()+".ttf",10); fp=n.lower()+".ttf"; break
        except IOError:
            try: ImageFont.truetype(n,10); fp=n; break
            except IOError: continue
    if fp:
        try: return ImageFont.truetype(fp, safe_base_size)
        except IOError: logger.warning(f"Font '{fp}' load failed with size {safe_base_size}. Defaulting.")
    try: return ImageFont.load_default(size=safe_base_size)
    except TypeError: return ImageFont.load_default()

def set_pil_image_format_to_png(image:Image.Image)->Image.Image:
    buf=io.BytesIO(); image.save(buf,format='PNG'); buf.seek(0)
    reloaded=Image.open(buf); reloaded.format="PNG"; return reloaded

def _derive_key(pw:str,salt:bytes)->bytes:
    kdf=PBKDF2HMAC(algorithm=hashes.SHA256(),length=KEY_SIZE,salt=salt,iterations=PBKDF2_ITERATIONS)
    return kdf.derive(pw.encode('utf-8'))

def encrypt_data(data:bytes,pw:str)->bytes:
    s=os.urandom(SALT_SIZE);k=_derive_key(pw,s);a=AESGCM(k);n=os.urandom(NONCE_SIZE)
    ct=a.encrypt(n,data,None); return s+n+ct

def decrypt_data(payload:bytes,pw:str)->bytes:
    ml=SALT_SIZE+NONCE_SIZE+TAG_SIZE;
    if len(payload)<ml: raise ValueError("Payload too short.")
    s,n,ct_tag=payload[:SALT_SIZE],payload[SALT_SIZE:SALT_SIZE+NONCE_SIZE],payload[SALT_SIZE+NONCE_SIZE:]
    k=_derive_key(pw,s);a=AESGCM(k)
    try: return a.decrypt(n,ct_tag,None)
    except InvalidTag: raise ValueError("Decryption failed: Invalid password/corrupted data.")
    except Exception as e: logger.error(f"Decrypt error: {e}",exc_info=True); raise

def _d2b(d:bytes)->str: return ''.join(format(b,'08b') for b in d)
def _b2B(b:str)->bytes:
    if len(b)%8!=0: raise ValueError("Bits not multiple of 8.")
    return bytes(int(b[i:i+8],2) for i in range(0,len(b),8))

def embed_data_in_image(img_obj:Image.Image,data:bytes)->Image.Image:
    img=img_obj.convert("RGB");px=np.array(img);fpx=px.ravel()
    lb=struct.pack('>I',len(data));fp=lb+data;db=_d2b(fp);nb=len(db)
    if nb>len(fpx): raise ValueError(f"Data too large: {nb} bits needed, {len(fpx)} available.")
    for i in range(nb): fpx[i]=(fpx[i]&0xFE)|int(db[i])
    spx=fpx.reshape(px.shape); return Image.fromarray(spx.astype(np.uint8),'RGB')

def extract_data_from_image(img_obj:Image.Image)->bytes:
    img=img_obj.convert("RGB");px=np.array(img);fpx=px.ravel()
    hbc=LENGTH_HEADER_SIZE*8
    if len(fpx)<hbc: raise ValueError("Image too small for header.")
    lb="".join(str(fpx[i]&1) for i in range(hbc))
    try: pl=struct.unpack('>I',_b2B(lb))[0]
    except Exception as e: raise ValueError(f"Header decode error: {e}")
    if pl==0: return b""
    if pl>(len(fpx)-hbc)/8: raise ValueError("Header len corrupted or > capacity.")
    tpb=pl*8; so=hbc; eo=so+tpb
    if len(fpx)<eo: raise ValueError("Image truncated or header corrupted.")
    pb="".join(str(fpx[i]&1) for i in range(so,eo)); return _b2B(pb)

def parse_kv_string_to_dict(kv_str:str)->dict:
    if not kv_str or not kv_str.strip(): return {}
    dd={};
    for ln,ol in enumerate(kv_str.splitlines(),1):
        l=ol.strip()
        if not l or l.startswith('#'): continue
        lc=l.split('#',1)[0].strip();
        if not lc: continue
        p=lc.split('=',1) if '=' in lc else lc.split(':',1) if ':' in lc else []
        if len(p)!=2: raise ValueError(f"L{ln}: Invalid format '{ol}'.")
        k,v=p[0].strip(),p[1].strip()
        if not k: raise ValueError(f"L{ln}: Empty key in '{ol}'.")
        dd[k]=v
    return dd

def convert_kb_to_kv_string(rules: list[str], memories: list[dict], include_rules: bool, include_memories: bool) -> str:
    lines = ["# iLearn Knowledge Base Export", f"# Exported on: {datetime.utcnow().isoformat()}Z"]
    
    if include_rules:
        lines.append("\n# --- RULES ---")
        for i, rule_text in enumerate(rules):
            lines.append(f"rule_{i+1} = {json.dumps(rule_text)}")

    if include_memories:
        lines.append("\n# --- MEMORIES ---")
        for i, mem_dict in enumerate(memories):
            lines.append(f"memory_{i+1} = {json.dumps(mem_dict)}")

    return "\n".join(lines)


def generate_brain_carrier_image(w=800, h=800) -> Image.Image:
    center_x, center_y = w / 2, h / 2
    y_coords, x_coords = np.mgrid[0:h, 0:w]

    distance = np.sqrt((x_coords - center_x)**2 + (y_coords - center_y)**2)
    max_distance = np.sqrt(center_x**2 + center_y**2)

    distance_norm = distance / max_distance

    bg_center_color = np.array([20, 25, 40])
    bg_outer_color = np.array([0, 0, 0])

    gradient = bg_outer_color + (bg_center_color - bg_outer_color) * (1 - distance_norm[..., np.newaxis])

    img = Image.fromarray(gradient.astype(np.uint8), 'RGB')
    draw = ImageDraw.Draw(img)

    num_distant_stars = int((w * h) / 200)
    for _ in range(num_distant_stars):
        x, y = random.randint(0, w - 1), random.randint(0, h - 1)
        brightness = random.randint(30, 90)
        draw.point((x, y), fill=(brightness, brightness, int(brightness * 1.1)))

    num_main_stars = int((w * h) / 1000)
    star_colors = [
        (255, 255, 255),
        (220, 230, 255),
        (255, 240, 220),
    ]

    for _ in range(num_main_stars):
        x, y = random.randint(0, w - 1), random.randint(0, h - 1)
        dist_from_center = np.sqrt((x - center_x)**2 + (y - center_y)**2)
        dist_ratio = min(dist_from_center / max_distance, 1.0)

        size = 0.5 + (2.5 * (dist_ratio ** 2))
        brightness = 120 + (135 * (dist_ratio ** 1.5))

        color = random.choice(star_colors)

        final_color = tuple(int(c * (brightness / 255.0)) for c in color)

        glow_size = size * 3
        glow_color = tuple(int(c * 0.3) for c in final_color)
        draw.ellipse([x - glow_size, y - glow_size, x + glow_size, y + glow_size], fill=glow_color)

        if random.random() < 0.15:
            draw.line([x-size, y, x+size, y], fill=final_color, width=1)
            draw.line([x, y-size, x, y+size], fill=final_color, width=1)
        else:
            draw.ellipse([x - size, y - size, x + size, y + size], fill=final_color)

    return img


def _get_text_measurement(draw_obj, text_str, font_obj):
    if hasattr(draw_obj, 'textbbox'):
        try:
            bbox = draw_obj.textbbox((0, 0), text_str, font=font_obj)
            width = bbox[2] - bbox[0]
            height = bbox[3] - bbox[1]
            return width, height
        except Exception: pass
    try:
        if hasattr(font_obj, 'getsize'): return font_obj.getsize(text_str)
        width, height = draw_obj.textsize(text_str, font=font_obj)
        return width, height
    except AttributeError:
        try:
            char_width_approx = font_obj.size * 0.6
            char_height_approx = font_obj.size
            return int(len(text_str) * char_width_approx), int(char_height_approx)
        except: return len(text_str) * 8, 10

def draw_key_list_dropdown_overlay(image: Image.Image, keys: list[str] = None, title: str = "Data Embedded") -> Image.Image:
    img_overlayed = image.copy().convert("RGBA")
    draw = ImageDraw.Draw(img_overlayed, "RGBA")
    width, height = img_overlayed.size

    overlay_color = (15, 23, 42, 190)
    title_color = (226, 232, 240)
    key_color = (148, 163, 184)

    font_bold = _get_font(PREFERRED_FONTS, 30)
    font_regular = _get_font(PREFERRED_FONTS, 15)

    draw.rectangle([0, 20, width, 80], fill=overlay_color)
    draw.text((width / 2, 50), title, fill=title_color, font=font_bold, anchor="ms")

    if keys:
        box_padding = 15
        line_spacing = 6
        text_start_x = 35
        lines = keys

        line_heights = [_get_text_measurement(draw, line, font_regular)[1] for line in lines]
        total_text_height = sum(line_heights) + (len(lines) - 1) * line_spacing
        box_height = total_text_height + (box_padding * 2)
        box_y0 = height - box_height - 20

        draw.rectangle([20, box_y0, width - 20, height - 20], fill=overlay_color)
        current_y = box_y0 + box_padding

        for i, key_text in enumerate(lines):
            draw.text((text_start_x, current_y), key_text, fill=key_color, font=font_regular)
            if i < len(line_heights):
                current_y += line_heights[i] + line_spacing

    final_image_rgb = Image.new("RGB", img_overlayed.size, (0, 0, 0))
    final_image_rgb.paste(img_overlayed, (0, 0), img_overlayed)

    return final_image_rgb