abyssallock.py

""" 
CLI tool for encrypting and decrypting files using the AbyssalLock engine.

python abyssallock.py encrypt <input_file> -p <passphrase>
python abyssallock.py decrypt <input_file> -p <passphrase>
"""

import hashlib, struct, os, argparse

# -----KDF: Key Derrived Function-----
def _kdf(seed: bytes, length: int) -> bytes:
    '''Algo will produce 32 bytes'''
    result = bytearray()
    counter = 0
    while len(result) < length:
        chunk = hashlib.sha256(seed + struct.pack('>I', counter)).digest() # Return 32 bytes
        result.extend(chunk)
        counter += 1
    return bytes(result[:length])

# -----Layer Base Class-----
class Layer:
    name: str = "Base Layer"
    description: str = ""
    
    def encrypt(self, data: bytes, key_bytes: bytes) -> bytes:
        raise NotImplementedError(f"{self.name}.encrypt() not implemented.")
    
    def decrypt(self, data: bytes, key_bytes: bytes) -> bytes:        
        raise NotImplementedError(f"{self.name}.decrypt() not implemented.")

# -----XOR Layer-----
class XORStreamLayer(Layer):
    name = "XOR Stream Layer"
    description = "Byte-level XOR against a key-derived pseudorandom stream"
    
    def encrypt(self, data: bytes, key_bytes: bytes) -> bytes:
        stream = _kdf(key_bytes, len(data))
        return bytes(a ^ b for a, b in zip(data, stream))
    
    def decrypt(self, data: bytes, key_bytes: bytes) -> bytes:
        return self.encrypt(data, key_bytes)

# -----S-Block Substitution Layer-----
class SubstitutionLayer(Layer):
    name = "S-Box Substitution"
    description = "RC4-style key-scheduled substitution table"
    
    def _build_sbox(self, key_bytes: bytes) -> list:
        sbox = list(range(256))
        j = 0
        for i in range(256):
            j = (j + sbox[i] + key_bytes[i % len(key_bytes)]) % 256
            sbox[i], sbox[j] = sbox[j], sbox[i]
        return sbox
    
    def encrypt(self, data: bytes, key_bytes: bytes) -> bytes:
        sbox = self._build_sbox(key_bytes)
        # look up each byte in data through sbox
        return bytes(sbox[byte] for byte in data)
    
    def decrypt(self, data: bytes, key_bytes: bytes) -> bytes:
        sbox = self._build_sbox(key_bytes)
        # build the inverse sbox, then look up each byte
        # hint: if sbox[3] = 17, then inv_sbox[17] = 3
        inv_sbox = [0] * 256
        for i, val in enumerate(sbox):
            inv_sbox[val] = i
        return bytes(inv_sbox[byte] for byte in data)

# -----Block Transposition Layer-----
class BlockTransposition(Layer):
    name = "Block Transposition Layer"
    description = "Fischer-Yates block shuffle based on a key-derived permutation"
    BLOCK_SIZE = 16
    
    def _permutation(self, key_bytes: bytes) -> list:
        bs = self.BLOCK_SIZE
        perm = list(range(bs))
        for i in range(bs - 1, 0, -1):
            j = key_bytes[i % len(key_bytes)] % (i + 1) # Pseudo random index
            perm[i], perm[j] = perm[j], perm[i]
        return perm
    
    def encrypt(self, data: bytes, key_bytes: bytes) -> bytes:
        bs = self.BLOCK_SIZE
        result = bytearray(data)
        for start in range(0, len(data) - bs + 1, bs):
            block = list(data[start: start + bs])
            perm = self._permutation(key_bytes)
            shuffeled = bytes(block[perm[j]] for j in range(bs))
            result[start: start + bs] = shuffeled
        return bytes(result)
    
    def decrypt(self, data: bytes, key_bytes: bytes) -> bytes:
        bs = self.BLOCK_SIZE
        perm = self._permutation(key_bytes)
        inv_perm = [0] * bs
        for i, val in enumerate(perm):
            inv_perm[val] = i
        result = bytearray(data)
        for i in range(0, len(data) - bs + 1, bs):
            block = data[i: i + self.BLOCK_SIZE]
            unshuffeled = bytes(block[inv_perm[j]] for j in range(bs))
            result[i: i + bs] = unshuffeled
        return bytes(result)

# -----Bit Rotation Layer-----
class BitRotation(Layer):
    name = "Bit Rotation Layer"
    description = "Per-byte bit rotation, rotation amount keyed per position"
    
    @staticmethod
    def _rotate_left(byte: int, shift: int) -> int:
        shift = shift % 8
        return ((byte << shift) | (byte >> (8 - shift))) & 0xFF
        
    @staticmethod
    def _rotate_right(byte: int, shift: int) -> int:
        shift = shift % 8
        return ((byte >> shift) | (byte << (8 - shift))) & 0xFF
    
    def encrypt(self, data: bytes, key_bytes: bytes) -> bytes:
        return bytes(self._rotate_left(byte, key_bytes[i % len(key_bytes)]) for i, byte in enumerate(data))
    
    def decrypt(self, data: bytes, key_bytes: bytes) -> bytes:
        return bytes(self._rotate_right(byte, key_bytes[i % len(key_bytes)]) for i, byte in enumerate(data))

# -----AbyssalLock Engine-----
class AbyssalLock:
    SALT_SIZE = 16
    KEY_SIZE = 512
    
    def __init__(self, passphrase: str):
        self.passphrase = passphrase
        self.layers = [XORStreamLayer(), SubstitutionLayer(), BlockTransposition(), BitRotation()]
        
    def _layer_key(self, layer_index: int, salt: bytes) -> bytes:
        seed = hashlib.sha256(
            f"{self.passphrase}:layer{layer_index}".encode() + salt
        ).digest()
        return _kdf(seed, 512)
    
    def encrypt(self, plaintext: bytes) -> bytes:
        salt = os.urandom(self.SALT_SIZE)
        data = plaintext
        for i, layer in enumerate(self.layers):
            key = self._layer_key(i, salt)
            data = layer.encrypt(data, key)
        return salt + data
    
    def decrypt(self, cyphertext: bytes) -> bytes:
        salt = cyphertext[:self.SALT_SIZE]
        data = cyphertext[self.SALT_SIZE:]
        for i, layer in reversed(list(enumerate(self.layers))):
            key = self._layer_key(i, salt)
            data = layer.decrypt(data, key)
        return data

def encrypt_file(path: str, passphrase: str) -> None:
    with open(path, "rb") as f:
        data = f.read()
    ct = AbyssalLock(passphrase).encrypt(data)
    out = path + ".enc"
    with open(out, "wb") as f:
        f.write(ct)
    print(f"[✓] Encrypted {path}→ {out}")

def decrypt_file(path: str, passphrase: str) -> None:
    with open(path, "rb") as f:
        data = f.read()
    pt = AbyssalLock(passphrase).decrypt(data)
    out = path.replace(".enc", "") if path.endswith(".enc") else path + ".dec"
    with open(out, "wb") as f:
        f.write(pt)
    print(f"[✓] Decrypted {path} → {out}")
    
def main():
    parser = argparse.ArgumentParser(description = "AbyssalLock Encryption Tool")
    parser.add_argument("action", choices = ["encrypt", "decrypt"])  
    parser.add_argument("file")
    parser.add_argument("-p", "--passphrase")
    args = parser.parse_args()
    
    if args.action == "encrypt":
        encrypt_file(args.file, args.passphrase)
    else:
        decrypt_file(args.file, args.passphrase)
        
def run_tests():
    out = _kdf(b"test-seed", 64)
    print(len(out))        # should print 64
    print(out.hex())       # should be deterministic — same every run
    
    out2 = _kdf(b"test-seed", 64)
    print(out == out2)     # should print True
    
    out3 = _kdf(b"different-seed", 64)
    print(out == out3)     # should print False
    
    layer = XORStreamLayer()
    ct = layer.encrypt(b"Hello", b"somekey")
    pt = layer.decrypt(ct, b"somekey")
    print(pt)  # should print b'Hello'
    
    sub = SubstitutionLayer()
    ct = sub.encrypt(b"Hello", b"somekey")
    pt = sub.decrypt(ct, b"somekey")
    print(pt)  # should print b'Hello'
    print(ct != b"Hello")  # should print True — it actually changed
    
    trans = BlockTransposition()
    ct = trans.encrypt(b"ABCDEFGHIJKLMNOP", b"somekey")
    pt = trans.decrypt(ct, b"somekey")
    print(pt)   # should print b'ABCDEFGHIJKLMNOP'
    print(ct != b"ABCDEFGHIJKLMNOP")  # should print True
    
    rot = BitRotation()
    ct = rot.encrypt(b"Hello", b"somekey")
    pt = rot.decrypt(ct, b"somekey")
    print(pt)  # should print b'Hello'
    print(ct != b"Hello")  # should print True
    
    engine = AbyssalLock("my-secret-passphrase")
    ct = engine.encrypt(b"Attack at dawn.")
    pt = engine.decrypt(ct)
    print(pt)  # should print b'Attack at dawn.'

    ct2 = engine.encrypt(b"Attack at dawn.")
    print(ct != ct2)  # should print True — different salt each time

# -----Verification Tests-----
if __name__ == "__main__":
    main()
    # run_tests()