# -*- coding: utf-8 -*-

import random as rnd

N = 64
R = 12
K = 64

def op_rotl(block, idx):
    return block[idx:]+block[:idx]

def op_xor(a, b):
    return bin(int(a,2)^int(b,2))[2:].zfill(N)

def op_perm(block, perm):
    return ''.join([block[perm[i]]for i in range(N)])

def op_comb(block):
    masked = bin(int('01' * (N//2),2) & int(block,2))[2:].zfill(N)
    return op_xor(op_rotl(masked, 1), block)

def inv_perm(perm):
    iperm = [-1] * N
    for i in range(N):
        iperm[perm[i]] = i
    return iperm

def encrypt(key, block, perm1, perm2):
    ct = op_xor(key, block)
    subkey = op_rotl(key, 3)
    for _ in range(R):
        ct = op_xor(ct, '11' + '0' * (N-2))
        ct = op_perm(ct, perm1)
        ct = op_comb(ct)
        ct = op_perm(ct, perm2)
        ct = op_comb(ct)
        ct = op_xor(subkey, ct)        
        subkey = op_rotl(subkey, 3)
    return ct

def decrypt(key, block, perm1, perm2):
    iperm1 = inv_perm(perm1)
    iperm2 = inv_perm(perm2)
    subkey = op_rotl(key, (3*R)%len(key))
    pt = block
    for _ in range(R):
        pt = op_xor(subkey, pt)
        subkey = op_rotl(subkey, -3)
        pt = op_comb(pt)
        pt = op_perm(pt, iperm2)
        pt = op_comb(pt)
        pt = op_perm(pt, iperm1)
        pt = op_xor(pt, '11' + '0' * (N-2))
    pt = op_xor(subkey, pt)
    return pt

def test():
    key = bin(rnd.getrandbits(N))[2:].zfill(N)
    pt = bin(rnd.getrandbits(N))[2:].zfill(N) 
    perm1 = list(range(N))
    rnd.shuffle(perm1)
    perm2 = list(range(N))
    rnd.shuffle(perm2)
    print(f'key: {int(key,2):x}')
    print(f'plaintext: {int(pt,2):x}')
    ct = encrypt(key, pt, perm1, perm2)
    print(f'ciphertext: {int(ct,2):x}')
    pt2 = decrypt(key, ct, perm1, perm2)
    print(f'decrypted: {int(pt2,2):x}')
    print(f'equal?: {pt2 == pt}')
    return


test()