"""
@file utils/encryption.py
@brief String encryption utilities for OMG-Fuscator.
@details Provides the StringEncryptor used by the obfuscation pipeline to transform
         string literals into runtime-decrypted expressions, making static analysis
         significantly harder.
"""

import base64
import hashlib
import random
from typing import Tuple

class StringEncryptor:
    """
    @brief Encrypts strings for obfuscation using layered XOR and Base85 encoding.
    @details Produces both an encoded payload and Python code that reconstructs
             the keys at runtime without embedding raw key bytes.
    """
    def __init__(self, primary_key: bytes, secondary_key: bytes, salt: bytes):
        """
        @brief Initialize encryptor keys.
        @param primary_key Primary XOR key bytes.
        @param secondary_key Secondary XOR key bytes.
        @param salt Salt bytes used to derive a per-string modifier.
        """
        self.primary_key = primary_key
        self.secondary_key = secondary_key
        self.salt = salt

    def hide_byte(self, b: int) -> str:
        """
        @brief Obfuscate a single key byte into an arithmetic/bitwise expression.
        @param b Input byte value in range [0, 255].
        @return Python expression string that evaluates to the original byte at runtime.
        """
        pattern = random.randint(1, 5)
        if pattern == 1:
            x1, x2, x3 = [random.randint(65, 90) for _ in range(3)]
            result = x1 ^ x2 ^ x3 ^ b
            return f"({x1}^{x2}^{x3}^{result})"
        elif pattern == 2:
            base = random.randint(65, 90)
            multiplier = random.randint(2, 4)
            adjusted = (base * multiplier - b)
            return f"(({base}*{multiplier}-{adjusted}))"
        elif pattern == 3:
            shift = random.randint(1, 3)
            modified = (b << shift) >> shift
            diff = b - modified
            return f"((({b << shift})>>{shift})+{diff})"
        elif pattern == 4:
            char1, char2 = random.sample(range(65, 91), 2)
            result = char1 + char2 - b
            return f"({char1}+{char2}-{result})"
        else:
            mod_base = random.randint(91, 100)
            result = b % mod_base
            factor = b // mod_base
            return f"({factor}*{mod_base}+{result})"

    def encrypt_string(self, s: str) -> Tuple[str, str, str]:
        """
        @brief Encrypt a string with layered XOR and Base85 encoding.
        @param s UTF-8 string to encrypt.
        @return Tuple[str, str, str]:
            - encoded: Base85 text of the encrypted payload
            - key_setup: Python code that reconstructs keys at runtime
            - modifier_hex: Hex string for the per-string modifier
        """
        data = s.encode('utf-8')
        modifier = hashlib.sha256(self.salt + data).digest()[:8]

        layer1 = bytes(
            d ^ k ^ m
            for d, k, m in zip(
                data,
                self.primary_key * ((len(data) // len(self.primary_key)) + 1),
                modifier * ((len(data) // len(modifier)) + 1)
            )
        )

        final_data = bytes(
            d ^ k
            for d, k in zip(
                layer1,
                self.secondary_key * ((len(layer1) // len(self.secondary_key)) + 1)
            )
        )

        encoded = base64.b85encode(final_data).decode()

        primary_key_code = f"bytes([{','.join(self.hide_byte(b) for b in self.primary_key)}])"
        secondary_key_code = f"bytes([{','.join(self.hide_byte(b) for b in self.secondary_key)}])"
        salt_code = f"bytes([{','.join(self.hide_byte(b) for b in self.salt)}])"

        key_setup = (
            f"_pk = {primary_key_code}\n"
            f"_sk = {secondary_key_code}\n"
            f"_st = {salt_code}"
        )
        
        return encoded, key_setup, modifier.hex()