import math
import random

MOD_ALL = 1114112

# Columnar Transposition

def columnar_encrypt(text, pad):
    cols = random.randint(1, len(text) - 1)
    rows = math.ceil(len(text) / cols)
    padded = text.ljust((rows * cols), str(chr(pad)))

    table = [padded[i:i+cols] for i in range(0, len(padded), cols)]
    result = ""
    for y in range(len(table[0])):
        for r in range(len(table)):
            result += table[r][y]
    return result, rows

def columnar_decrypt(ciphertext, key, pad):
    rows = key
    table = [ciphertext[i:i+rows] for i in range(0, len(ciphertext), rows)]
    
    res = ""
    for a in range(rows):
        for b in range(len(table)):
            res += table[b][a]
            
    return res.rstrip(str(chr(pad)))

# Hill Encrypt

def keygen(l):
    while True:
        if l % 2 == 0:
            a, b, c, d = [random.randint(0, MOD_ALL) for _ in range(4)]
            det = (a*d - b*c) % MOD_ALL
            if det < MOD_ALL and math.gcd(det, MOD_ALL) == 1:
                return f"{a}-{b}-{c}-{d}"
        else:
            a, b, c, d, e, f, g, h ,i = [random.randint(0, 99) for _ in range(9)]
            det = (a*e*i + b*f*g + c*d*h - b*d*i - a*f*h - c*e*g) % MOD_ALL
            if det < MOD_ALL and math.gcd(det, MOD_ALL) == 1:
                return f"{a}-{b}-{c}-{d}-{e}-{f}-{g}-{h}-{i}"

def key_separator(k):
    return k.split("-")

def hexcon(h):
    return str(format(h, "06X"))

def hill_encrypt(pl, k):
    cipher = ""
    if len(pl) % 2 == 0:
        for i in range(0, len(pl), 2):
            vec = [int(ord(pl[i])), int(ord(pl[i+1]))]
            bl1 = (int(k[0])*vec[0] + int(k[1])*vec[1]) % MOD_ALL
            bl2 = (int(k[2])*vec[0] + int(k[3])*vec[1]) % MOD_ALL
            cipher += hexcon(bl1) + hexcon(bl2)
    else:
        for i in range(0, len(pl), 3):
            vec = [int(ord(pl[i])), int(ord(pl[i+1])), int(ord(pl[i+2]))]
            bl1 = (int(k[0])*vec[0] + int(k[1])*vec[1] + int(k[2])*vec[2]) % MOD_ALL
            bl2 = (int(k[3])*vec[0] + int(k[4])*vec[1] + int(k[5])*vec[2]) % MOD_ALL
            bl3 = (int(k[6])*vec[0] + int(k[7])*vec[1] + int(k[8])*vec[2]) % MOD_ALL
            cipher += hexcon(bl1) + hexcon(bl2) + hexcon(bl3)
    return cipher

def matrix_inverse_2x2(k):
    a, b, c, d = map(int, k)
    det = (a*d - b*c) % MOD_ALL
    inv_det = pow(det, -1, MOD_ALL)
    return [( d*inv_det) % MOD_ALL,
            (-b*inv_det) % MOD_ALL,
            (-c*inv_det) % MOD_ALL,
            ( a*inv_det) % MOD_ALL]

def matrix_inverse_3x3(k):
    a,b,c,d,e,f,g,h,i = map(int, k)
    det = (a*e*i + b*f*g + c*d*h - b*d*i - a*f*h - c*e*g) % MOD_ALL
    inv_det = pow(det, -1, MOD_ALL)
    A = (e*i - f*h) % MOD_ALL
    B = (-(b*i - c*h)) % MOD_ALL
    C = (b*f - c*e) % MOD_ALL
    D = (-(d*i - f*g)) % MOD_ALL
    E = (a*i - c*g) % MOD_ALL
    F = (-(a*f - c*d)) % MOD_ALL
    G = (d*h - e*g) % MOD_ALL
    H = (-(a*h - b*g)) % MOD_ALL
    I = (a*e - b*d) % MOD_ALL
    adj = [A, B, C,
           D, E, F,
           G, H, I]
    return [(x * inv_det) % MOD_ALL for x in adj]

def hex_to_int_blocks(cipher, block_size):
    values = [int(cipher[i:i+6], 16) for i in range(0, len(cipher), 6)]
    return [values[i:i+block_size] for i in range(0, len(values), block_size)]

def hill_decrypt(cipher, k, pad):
    plain = ""
    if len(k) == 4:
        invK = matrix_inverse_2x2(k)
        blocks = hex_to_int_blocks(cipher, 2)
        for vec in blocks:
            x = (invK[0]*vec[0] + invK[1]*vec[1]) % MOD_ALL
            y = (invK[2]*vec[0] + invK[3]*vec[1]) % MOD_ALL
            plain += chr(x) + chr(y)
    else:
        invK = matrix_inverse_3x3(k)
        blocks = hex_to_int_blocks(cipher, 3)
        for vec in blocks:
            x = (invK[0]*vec[0] + invK[1]*vec[1] + invK[2]*vec[2]) % MOD_ALL
            y = (invK[3]*vec[0] + invK[4]*vec[1] + invK[5]*vec[2]) % MOD_ALL
            z = (invK[6]*vec[0] + invK[7]*vec[1] + invK[8]*vec[2]) % MOD_ALL
            plain += chr(x) + chr(y) + chr(z)
    return plain.rstrip(str(chr(pad)))

if __name__ == "__main__":
    print("="*100)
    print("Welcome to Hill Fences Superencryption")
    print("1.Encryption\n2.Decryption\n")
    opt = int(input("Select an option: "))
    print("="*100)
    if opt == 1:
        padding_random_hill = random.randint(0, MOD_ALL)
        padding_random_columnar = random.randint(0, MOD_ALL)
        text = input("Write your text: \n")
        col_encrypted = columnar_encrypt(text, padding_random_columnar)
        col_res, key = col_encrypted
        padding_random_ = random.randint(0, MOD_ALL)
        if len(col_res) % 2 != 0 and len(col_res) % 3 != 0:
            col_res += str(chr(padding_random_hill))
        hill_key_str = keygen(len(col_res))
        hill_key = key_separator(hill_key_str)
        final = hill_encrypt(col_res, hill_key)
        print("\n"+ "="*100)
        print("Final Cipher:\n" + final + "\n")
        print(f"Key: \n{padding_random_columnar}-{padding_random_hill}-{key}-{hill_key_str}")
    elif opt == 2:
        text = input("Write your cipher: \n")
        print()
        key = input("Write your key: ")
        splitted_key = key.split("-")
        hill_key = [int(x) for x in splitted_key[3:]]
        hill_decrypted = hill_decrypt(text, hill_key, int(splitted_key[1]))
        final_plain = columnar_decrypt(hill_decrypted, int(splitted_key[2]), int(splitted_key[0]))
        print()
        print("Decrypted message: \n" + final_plain)
    print("\n"+ "="*100)