# solve.py
import numpy as np

# Load tables
matrix_raw = open("/tmp/matrix_1253.bin", "rb").read()
table_a = open("/tmp/table_2605.bin", "rb").read()  # 55 bytes
table_b = open("/tmp/table_2697.bin", "rb").read()  # 55 bytes
expected_raw = open("/tmp/expected_2119.bin", "rb").read()  # 110 bytes

# Build coefficient matrix mod 256
# coeff[i][j] = ((table_a[i] * table_b[j] * 1337) & 0xFF + matrix[i*55+j]) & 0xFF
M = np.zeros((55, 55), dtype=np.int64)
for i in range(55):
    for j in range(55):
        ab = table_a[i] * table_b[j]       # 8x8 -> 16-bit
        val = (ab * 1337) & 0xFF            # imul then take low byte
        val = (val + matrix_raw[i * 55 + j]) & 0xFF  # add matrix element
        M[i][j] = val

# Expected values — try byte extraction (every other byte = low bytes of words)
# The cmpsw comparison uses words, so expected is 55 x 2-byte values
# Result bytes are at rsp[0..54], compared as words: (rsp[0],rsp[1]), (rsp[2],rsp[3])...
# Since only odd-indexed bytes have values, let's try both interpretations

# Interpretation 1: expected is just the first 55 bytes
expected_bytes = list(expected_raw[:55])

# Interpretation 2: expected is every other byte (word low bytes)  
expected_words_lo = [expected_raw[i*2] for i in range(55)]

print("Expected (first 55 bytes):", [hex(x) for x in expected_bytes[:10]])
print("Expected (word lo bytes):", [hex(x) for x in expected_words_lo[:10]])

# Gaussian elimination mod 256
def solve_mod256(M, b):
    n = len(b)
    # Augmented matrix
    aug = np.zeros((n, n+1), dtype=np.int64)
    for i in range(n):
        for j in range(n):
            aug[i][j] = M[i][j] % 256
        aug[i][n] = b[i] % 256

    # Forward elimination
    for col in range(n):
        # Find pivot with odd value (invertible mod 256 requires gcd(pivot,256)=1)
        pivot = -1
        for row in range(col, n):
            if aug[row][col] % 2 == 1:  # odd = invertible mod 256
                pivot = row
                break
        if pivot == -1:
            print(f"No odd pivot at column {col}, trying any nonzero...")
            for row in range(col, n):
                if aug[row][col] != 0:
                    pivot = row
                    break
            if pivot == -1:
                print(f"Singular at column {col}")
                return None

        # Swap
        aug[[col, pivot]] = aug[[pivot, col]]

        # Compute inverse of pivot mod 256
        pv = int(aug[col][col]) % 256
        inv = pow(pv, -1, 256) if pv % 2 == 1 else None
        if inv is None:
            print(f"Non-invertible pivot {pv} at col {col}")
            return None

        # Scale pivot row
        aug[col] = (aug[col] * inv) % 256

        # Eliminate
        for row in range(n):
            if row != col and aug[row][col] != 0:
                factor = int(aug[row][col])
                aug[row] = (aug[row] - factor * aug[col]) % 256

    solution = aug[:, n] % 256
    return solution

# Try both interpretations
for name, expected in [("first_55_bytes", expected_bytes),
                        ("word_lo_bytes", expected_words_lo)]:
    print(f"\n=== Trying {name} ===")
    sol = solve_mod256(M, expected)
    if sol is not None:
        flag_content = bytes([int(x) for x in sol])
        flag = b"CTF{th1s" + flag_content + b"}"
        printable = all(0x20 <= b < 0x7f for b in flag_content)
        print(f"Solution: {flag}")
        print(f"Printable: {printable}")
        print(f"Hex: {flag_content.hex()}")