HackTheBox Permuted Challenge

#!/usr/bin/env python3
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from Crypto.Util.number import long_to_bytes
from hashlib import sha256
from math import gcd
from functools import reduce

class Permutation:
    def __init__(self, mapping):
        if isinstance(mapping, Permutation):
            self.mapping = tuple(mapping.mapping)
        else:
            self.mapping = tuple(mapping)
        self.length = len(self.mapping)

    def __call__(self, x):
        return self.mapping[x]

    def __mul__(self, other):
        return Permutation([self(other(i)) for i in range(self.length)])

    def __pow__(self, n):
        if n == 0:
            return Permutation(range(self.length))
        if n == 1:
            return self
        if n < 0:
            raise ValueError("Negative powers not implemented")
        half = self ** (n // 2)
        if n % 2 == 0:
            return half * half
        return half * half * self

    def cycles(self):
        cycles = []
        used = set()
        for i in range(self.length):
            if i in used:
                continue
            current_cycle = [i]
            used.add(i)
            next_idx = self(i)
            while next_idx not in used:
                current_cycle.append(next_idx)
                used.add(next_idx)
                next_idx = self(next_idx)
            if len(current_cycle) > 1:
                cycles.append(current_cycle)
        return sorted(cycles)

def extended_gcd(a, b):
    if a == 0:
        return b, 0, 1
    gcd, x1, y1 = extended_gcd(b % a, a)
    x = y1 - (b // a) * x1
    y = x1
    return gcd, x, y

def lcm(a, b):
    return abs(a * b) // gcd(a, b)

def crt(moduli, remainders):
    if not moduli or not remainders:
        raise ValueError("Empty inputs to CRT")
    if len(moduli) != len(remainders):
        raise ValueError("Number of moduli must equal number of remainders")
    if len(moduli) == 1:
        return remainders[0] % moduli[0]
    result = remainders[0]
    modulus = moduli[0]
    for m2, r2 in zip(moduli[1:], remainders[1:]):
        g, p, q = extended_gcd(modulus, m2)
        if r2 % g != result % g:
            continue
        new_mod = (modulus * m2) // g
        result = (result + ((((r2 - result) * p * modulus) // g) % new_mod)) % new_mod
        modulus = new_mod
    return result

def solve_dlp(g, h):
    g_cycles = g.cycles()
    h_cycles = h.cycles()
    print(f"Number of g cycles: {len(g_cycles)}")
    print(f"Number of h cycles: {len(h_cycles)}")
    G = []
    H = []
    for i in range(g.length):
        g_pos = None
        h_pos = None
        for j, c in enumerate(g_cycles):
            if i in c:
                g_pos = (j, c.index(i))
        for j, c in enumerate(h_cycles):
            if i in c:
                h_pos = (j, c.index(i))
        if g_pos is not None:
            G.append(g_pos)
        if h_pos is not None:
            H.append(h_pos)
    First = []
    Second = []
    for c in h_cycles:
        if len(c) > 1:
            First.append(c[0])
            Second.append(c[1])
    D = []
    L = []
    for i in range(len(Second)):
        dist = G[Second[i]][1] - G[First[i]][1]
        cycle_len = len(h_cycles[i])
        if cycle_len > 1:
            D.append(dist % cycle_len)
            L.append(cycle_len)
    print("D:", D)
    print("L:", L)
    try:
        return crt(L, D)
    except Exception as e:
        print(f"CRT error: {e}")
        raise

def decrypt_shared_secret(B, a, ciphertext):
    try:
        C = B**a
        sec = tuple(C.mapping)
        sec = abs(hash(sec))
        sec = long_to_bytes(sec)
        hash_obj = sha256()
        hash_obj.update(sec)
        key = hash_obj.digest()[16:32]
        iv = b'\xfb\x1d]\xbc\r\xa3\x91\x1cvV#\x13\xd8z\xb4\x16'
        cipher = AES.new(key, AES.MODE_CBC, iv)
        decrypted = cipher.decrypt(ciphertext)
        try:
            return unpad(decrypted, 16)
        except ValueError:
            return decrypted
    except Exception as e:
        print(f"Decryption error: {e}")
        return None

def main():
    print("Loading challenge data...")
    namespace = {}
    with open('output.txt', 'r') as f:
        exec(f.read(), namespace)
    g = Permutation(namespace['g'])
    A = Permutation(namespace['A'])
    B = Permutation(namespace['B'])
    c = namespace.get('c', None)
    if not c:
        print("Error: No ciphertext found in output.txt")
        return
    print("Solving DLP...")
    try:
        a = solve_dlp(g, A)
        print(f"Found private key a = {a}")
        if g ** a == A:
            print("Verified: g^a = A")
        else:
            print("Warning: g^a != A")
        print("Decrypting flag...")
        flag = decrypt_shared_secret(B, a, c)
        if flag:
            print(flag.decode('ascii'))
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()