test_encryption.ts

import * as crypto from 'crypto';
import keysConfig from './keys.json';
import { 
    encryptSubstitution, 
    decryptSubstitution,
    encryptCaesar, 
    decryptCaesar,
    railFenceEncrypt, 
    railFenceDecrypt,
    addHiddenPadding, 
    removeHiddenPadding,
    getKeyShifts, 
    getRepeatingShifts,
    encryptMath,
    decryptMath,
    stretchKey,
    generateHMAC,
    generateRandomSalt,
    generateRandomIV,
    deriveKeys,
    aesEncrypt,
    aesDecrypt,
    getPipelineOrder
} from './cipher';

export function testFullRoundTrip(plaintext: string, key: string, numRounds: number): boolean {
    console.log(`\n--- Testing Round Trip ---`);
    console.log(`Plaintext Length: ${plaintext.length}`);
    console.log(`Key length: ${key.length}, Rounds: ${numRounds}`);

    // --- ENCRYPTION ---
    const liveSalt = generateRandomSalt();
    const masterKeyBuffer = stretchKey(key, liveSalt);
    const { aesKey, hmacKey } = deriveKeys(masterKeyBuffer, liveSalt);
    const masterKeyHex = masterKeyBuffer.toString('hex');

    let state = plaintext;
    let iv = "";
    let tag = "";
    let hmac = "";

    const masterPipeline = (keysConfig as any).MASTER_PIPELINE;
    const legacyKeyShifts = getKeyShifts(key);
    const legacyRepeatingShifts = getRepeatingShifts(legacyKeyShifts, numRounds);

    for (const stage of masterPipeline) {
        if (stage === "substitution" && (keysConfig as any).ENABLE_SUBSTITUTION_LAYER) {
            state = encryptSubstitution(state);
        } else if (stage === "chaos") {
            for (let round = 0; round < numRounds; round++) {
                const legacyRoundShift = legacyRepeatingShifts[round];
                const roundDigest = crypto.createHash('sha256').update(masterKeyHex + round.toString() + legacyRoundShift.toString()).digest();
                const shift = roundDigest[0]; 
                const modifier = roundDigest[1]; 
                const pipelineOrder = getPipelineOrder(modifier);

                for (const step of pipelineOrder) {
                    if (step === "caesar" && (keysConfig as any).ENABLE_CAESAR_LAYER) state = encryptCaesar(state, shift);
                    if (step === "substitution" && (keysConfig as any).ENABLE_SUBSTITUTION_LAYER) state = encryptSubstitution(state);
                    if (step === "rail_fence" && (keysConfig as any).ENABLE_RAIL_FENCE_LAYER) state = railFenceEncrypt(state);
                    if (step === "math" && (keysConfig as any).ENABLE_MATH_LAYER) state = encryptMath(state, shift);
                }
            }
        } else if (stage === "padding") {
            state = addHiddenPadding(state);
        } else if (stage === "aes" && (keysConfig as any).AES_ENABLED) {
            iv = generateRandomIV();
            const aesResult = aesEncrypt(state, aesKey, iv);
            state = aesResult.ciphertext;
            tag = aesResult.tag;
        } else if (stage === "hmac" && (keysConfig as any).ENABLE_HMAC) {
            const hmacPayload = state + liveSalt + iv + tag;
            hmac = generateHMAC(hmacKey, hmacPayload);
        }
    }

    // Final normalization to Base64 for transport.
    const lastStageForwardEn = masterPipeline[masterPipeline.length - 1];
    const finalData = (lastStageForwardEn === "aes" && (keysConfig as any).AES_ENABLED) ? state : Buffer.from(state, 'utf16le').toString('base64');

    const payloadObj = { data: finalData, hmac, salt: liveSalt, iv, tag };
    const jsonStr = JSON.stringify(payloadObj);
    const securePayload = Buffer.from(jsonStr, 'utf-8').toString('base64');

    console.log(`[Encrypted Payload Structure Generated] -> ${securePayload.length} bytes`);

    // --- DECRYPTION ---
    const payloadStrParsed = Buffer.from(securePayload, 'base64').toString('utf-8');
    const parsedObj = JSON.parse(payloadStrParsed);

    const dsalt = parsedObj.salt;
    const masterKeyBufferDec = stretchKey(key, dsalt);
    const { aesKey: aesKeyDec, hmacKey: hmacKeyDec } = deriveKeys(masterKeyBufferDec, dsalt);
    const masterKeyHexDec = masterKeyBufferDec.toString('hex');
    
    let dstate = parsedObj.data;
    let div = parsedObj.iv;
    let dtag = parsedObj.tag;
    let dhmac = parsedObj.hmac;

    // Determination: If the last stage (forward) was not AES, we must decode the transport layer first.
    if (lastStageForwardEn !== "aes" || !((keysConfig as any).AES_ENABLED)) {
        dstate = Buffer.from(dstate, 'base64').toString('utf16le');
    }

    for (let i = masterPipeline.length - 1; i >= 0; i--) {
        const stage = masterPipeline[i];
        
        if (stage === "hmac" && (keysConfig as any).ENABLE_HMAC) {
            const hmacPayload = dstate + dsalt + div + dtag;
            const calculatedHmac = generateHMAC(hmacKeyDec, hmacPayload);
            if (calculatedHmac !== dhmac) {
                console.error("❌ FAILURE: Authetication failed! HMAC mismatch.");
                return false;
            }
        } else if (stage === "aes" && (keysConfig as any).AES_ENABLED) {
            dstate = aesDecrypt(dstate, aesKeyDec, div, dtag);
        } else if (stage === "padding") {
            dstate = removeHiddenPadding(dstate);
        } else if (stage === "chaos") {
            let chaosText = dstate;
            for (let round = numRounds - 1; round >= 0; round--) {
                const legacyRoundShift = legacyRepeatingShifts[round];
                const roundDigest = crypto.createHash('sha256').update(masterKeyHexDec + round.toString() + legacyRoundShift.toString()).digest();
                const shift = roundDigest[0]; 
                const modifier = roundDigest[1]; 
                const pipelineOrder = getPipelineOrder(modifier);

                for (let s = pipelineOrder.length - 1; s >= 0; s--) {
                    const step = pipelineOrder[s];
                    if (step === "math" && (keysConfig as any).ENABLE_MATH_LAYER) chaosText = decryptMath(chaosText, shift);
                    if (step === "rail_fence" && (keysConfig as any).ENABLE_RAIL_FENCE_LAYER) chaosText = railFenceDecrypt(chaosText);
                    if (step === "substitution" && (keysConfig as any).ENABLE_SUBSTITUTION_LAYER) chaosText = decryptSubstitution(chaosText);
                    if (step === "caesar" && (keysConfig as any).ENABLE_CAESAR_LAYER) chaosText = decryptCaesar(chaosText, shift);
                }
            }
            dstate = chaosText;
        } else if (stage === "substitution" && (keysConfig as any).ENABLE_SUBSTITUTION_LAYER) {
            dstate = decryptSubstitution(dstate);
        }
    }

    const decryptedText = dstate;

    if (plaintext === decryptedText) {
        console.log("✅ SUCCESS: Decrypted text perfectly matches the original!");
        return true;
    } else {
        console.error("❌ FAILURE: Mismatch detected!");
        return false;
    }
}

function testAdversarialHMAC(): boolean {
    console.log(`\n--- Testing Adversarial Tampering (HMAC) ---`);
    const plaintext = "Top Secret Payload";
    const key = "adminKey";
    
    // Encrypt
    const liveSalt = generateRandomSalt();
    const mBuffer = stretchKey(key, liveSalt);
    const { hmacKey } = deriveKeys(mBuffer, liveSalt);
    const finalData = Buffer.from("fake data", 'utf-8').toString('base64');
    let hmac = generateHMAC(hmacKey, finalData + liveSalt + "" + "");
    
    // Tamper with payload
    const hackedData = Buffer.from("hacked data", 'utf-8').toString('base64');
    const payloadObj = { data: hackedData, hmac, salt: liveSalt, iv: "", tag: "" };
    
    console.log("Attempting to decrypt tampered payload...");
    const calculatedHmac = generateHMAC(hmacKey, payloadObj.data + payloadObj.salt + payloadObj.iv + payloadObj.tag);
    if (calculatedHmac !== payloadObj.hmac) {
        console.log("✅ SUCCESS: HMAC correctly blocked the tampered payload!");
        return true;
    } else {
        console.error("❌ FAILURE: System accepted corrupted payload!");
        return false;
    }
}

function testWrongKeyHMAC(): boolean {
    console.log(`\n--- Testing Adversarial Tampering (Wrong Key) ---`);
    const plaintext = "This message is for authorized eyes only.";
    const validKey = "authorized123";
    const hackerKey = "password123";

    const liveSalt = generateRandomSalt();
    const vBuffer = stretchKey(validKey, liveSalt);
    const { hmacKey: validHmacKey } = deriveKeys(vBuffer, liveSalt);
    const finalData = Buffer.from("secret base64 string", 'utf-8').toString('base64');
    let hmac = generateHMAC(validHmacKey, finalData + liveSalt + "" + "");
    
    const payloadObj = { data: finalData, hmac, salt: liveSalt, iv: "", tag: "" };
    
    console.log("Attempting to decrypt payload with WRONG KEY...");
    const hackerBuffer = stretchKey(hackerKey, payloadObj.salt);
    const { hmacKey: hackerHmacKey } = deriveKeys(hackerBuffer, payloadObj.salt);
    const calculatedHmac = generateHMAC(hackerHmacKey, payloadObj.data + payloadObj.salt + payloadObj.iv + payloadObj.tag);
    
    if (calculatedHmac !== payloadObj.hmac) {
        console.log("✅ SUCCESS: Decryption instantly denied. Wrong Key HMAC securely rejected payload!");
        return true;
    } else {
        console.error("❌ FAILURE: System accepted wrong key payload!");
        return false;
    }
}

// Only run tests if this file is executed directly
if (import.meta.url === `file://${process.argv[1]}` || process.argv[1]?.endsWith('test_encryption.ts')) {
    const testsPassed = [
        testFullRoundTrip("Hello World!", "secret123", 5),
        testFullRoundTrip("This is a long sentence to thoroughly test out the padding algorithm. It has punctuation!", "password!@#", 12),
        testFullRoundTrip("A", "single_char", 1),
        testFullRoundTrip("", "empty_string_test", 1),
        testFullRoundTrip("A".repeat(50), "long_text", 3),
        testFullRoundTrip("1a2b3c4d5e".repeat(10), "huge", 5),
        testFullRoundTrip("!@#$%^&*()_+{}|:\"<>?-=[]\\;',./", "symbols", 7),
        testFullRoundTrip("     ", "spaces_only", 2),
        testFullRoundTrip("newlines\nand\ttabs\r\0", "whitespace_key", 4),
        testFullRoundTrip("こんにちは世界！", "japanese_key", 3),
        testFullRoundTrip("Привет, мир!", "russian_key", 3),
        testFullRoundTrip("مرحباً بالعالم", "arabic_key", 3),
        testFullRoundTrip("🚀👨‍👩‍👧‍👦🔥", "emoji_key", 4),
        testFullRoundTrip("Test Zero Rounds", "zero", 0),
        testAdversarialHMAC(),
        testWrongKeyHMAC()
    ];

    if (testsPassed.every(Boolean)) {
        console.log("\n🎉 ALL TESTS PASSED SUCCESSFULLY! The Framework is fully functional.");
        process.exit(0);
    } else {
        console.error("\n⚠️ SOME TESTS FAILED. There is a flaw in the Cryptographic framework integration.");
        process.exit(1);
    }
}