# Chris M. Thomasson Copyright 2018 (c)
# Experimental HMAC Cipher
# Test Vector For C program at
# https://pastebin.com/raw/feUnA3kP
#____________________________________________________________


# Our external libs
#____________________________________________________________
import random;
import hashlib;
import hmac;


# Some Utilities
#____________________________________________________________
def ct_bytes_to_hex(origin, offset):
    hex = "";
    n = len(origin);
    t = "0123456789ABCDEF";
    for i in range(offset, n):
        c = ord(origin[i]);
        nibl = c & 0x0F;
        nibh = (c & 0xF0) >> 4;
        hex = hex + t[nibh];
        hex = hex + t[nibl];
        hex = hex + " ";
        if (not ((i + 1) % 16) and i != n - 1): 
            hex = hex + "\r\n";
    return hex;


# Generate n random bytes
# These need should ideally be from a truly random, non-repeatable 
# source. TRNG!
def ct_rand_bytes(n):
    rb = "";
    for i in range(n):
        #rb = rb + chr(random.randint(0, 255));
        # HACK to get the same numbers
        rb = rb + chr(i);
    return rb;


# The Secret Key
# Contains all the parts of the secret key
#____________________________________________________________
class ct_secret_key:
    def __init__(self, hmac_key, hash_algo, rand_n):
        self.hmac_key = hmac_key;
        self.hash_algo = hash_algo;
        self.rand_n = rand_n;
 
    def __repr__(self):
        return "hmac_key:%s\nhash_algo:%s\nrand_n:%s" % (ct_bytes_to_hex(self.hmac_key, 0), self.hash_algo, self.rand_n);
 
    def __str__(self): return self.__repr__();


# The Ciphertext or Plaintext
# It holds the bytes of a ciphertext or a plaintext
#____________________________________________________________
class ct_bin:
    def __init__(self, ctxt):
        self.bytes = ctxt;
    def __repr__(self):
        return "%s" % (ct_bytes_to_hex(self.bytes, 0));
 
    def __str__(self): return self.__repr__();


# The Crypt Round Function
#____________________________________________________________
def ct_crypt_round(SK, P, M):
    H = hmac.new(SK.hmac_key.encode(), None, SK.hash_algo);
    H.update(SK.hmac_key[::-1].encode());
    C = "";
    I_P = 0;
    I_P_N = len(P.bytes);

    while (I_P < I_P_N):
        D = H.digest();
        print("D:%s" % (H.hexdigest()));
        I_D = 0;
        I_D_N = len(D);
        while (I_P < I_P_N and I_D < len(D)):
            P_byte = ord(P.bytes[I_P]);
            C_I_P = P_byte ^ D[I_D];
            C = C + chr(C_I_P);
            if (M == False):
                U = bytearray([P_byte, C_I_P]);
                H.update(U);
            else:
                U = bytearray([C_I_P, P_byte]);
                H.update(U);
            I_P = I_P + 1;
            I_D = I_D + 1;
    return ct_bin(C);


# The Crypt Function
#____________________________________________________________
def ct_crypt(SK, P, M):
    if (M == False):
        R = ct_rand_bytes(SK.rand_n);
        P.bytes = R + P.bytes;
    C = ct_crypt_round(SK, P, M);
    C_1 = ct_bin(C.bytes[::-1]);
    C = ct_crypt_round(SK, C_1, M);
    if (M == True):
        size = len(C.bytes) - SK.rand_n;
        C.bytes = C.bytes[SK.rand_n : SK.rand_n + size];
    return C;


# The Main Program
#____________________________________________________________

# Alice and Bob's Secret Key
#____________________
SK = ct_secret_key(
    "Password",
    hashlib.sha256, # The hash function. It should be a crypto secure hash.
    32 # The number of bytes. The should be generated by a TRNG
);
print("%s" % (SK));

# Alice's Plaintext
#____________________
Original_Plaintext = "Plaintext";
A_P = ct_bin(Original_Plaintext);
print(
    "\n\nAlice's Plaintext Bytes:"
    "\n____________________\n%s\n" % (A_P)
);

# Encrypt
#____________________
C = ct_crypt(SK, A_P, False);
print(
    "\n\nCiphertext Bytes:"
    "\n____________________\n%s\n" % (C)
);

# Decrypt
#____________________
B_P = ct_crypt(SK, C, True);
print(
    "\n\nBob's Ciphertext Bytes:"
    "\n____________________\n%s\n" % (B_P)
);


if (B_P.bytes != Original_Plaintext):
  print("DATA CORRUPTED!");