#!/usr/bin/env python3

from random import choice
from string import ascii_uppercase
from string import digits

# Variables for Rail Fence Transposition Cipher
# depth of rail fence cipher - numbe rof rows
k = 3
clear_text_list = []
clear_text_list_element = ''
count_char = 1
i = 0
encrypted_text = ''

# Variables for ROT cipher
rot=5
cipher_text = ""

with open("clear_text.txt", "r") as f:
    # we need to remove last char
    clear_text = (f.read()[:-1]).upper()

#### Rail Fence Transposition Cipher ##### 

# Compute real text length 
text_length = max_text_length = len(clear_text)
print(f"Length of clear text with spaces is: {text_length}")

# Compute max length
while (max_text_length % k != 0):
    max_text_length+=1

# Compute difference between max and real lenght
rozdiel = max_text_length - text_length
print(f"For k = {k} we must add '{rozdiel}' random chars to '{text_length}' so result of multiplying of number k is equal or bigger than '{text_length}', so '{max_text_length}'")

# Generate number of lowecase char that match ivar rozdiel and them to end of clear_text
strings = ''.join(choice(ascii_uppercase) for i in range(rozdiel))
print(f"Those are randomly generated chars: '{strings}' that we add to end of clear text")

# Add strings to end of variable
clear_text = clear_text + strings
print(f"Length of clear text with added strings '{strings}' to clear text is: {len(clear_text)}")

# Display clear that with added chars that is clear text for rot cipher
print(f"\nClear text with added chars is input into ROT{rot} cipher: ")
print(clear_text, "\n")


#### ROT Cipher #####

# Generate variable ABC..XYZ0123..9
template = "(,;'<.$![@" + ascii_uppercase + '}>+:&-%]*\\ ' + digits + '_"?=`/~#|){'

print(f"Template for ROT{rot} is: {template}")

# Get length of template
len_template = len(template)
print(f"Length of template is: '{len_template}'")

for pismeno in clear_text:
    # get index of pismeno and add length of rot cipher
    # if pismeno not find in template, position is -1 so rot_position is 4 for ROT5 (-1 + 5)
    rot_position = template.find(pismeno) + rot
    print(f"Char is: {pismeno} and ROT position is: {rot_position}")

    if rot_position < len_template:
        cipher_text += template[rot_position]
    elif rot_position > len_template:
        cipher_text += template[rot_position-len_template]

# Display ROT cipher text
print(f"\nThis is ROT{rot} cipher text that is input to Rail Fence Transposition Cipher: ")
clear_text = cipher_text
print(clear_text, "\n")

#### Rail Fence Transposition Cipher ####

# Create list with k-number fields
for pismeno in clear_text:
        
    if count_char % k != 0:
        clear_text_list_element = clear_text_list_element + pismeno
    elif count_char % k == 0:
        clear_text_list_element = clear_text_list_element + pismeno
        clear_text_list.append(clear_text_list_element)
        clear_text_list_element = ''

    count_char+=1

print(f"ROT{rot} encrypted text added to list, each element is k = {k} length: ")
print(clear_text_list, "\n")

# Read first char from list element than move to another element and read first char until end of list
# then read second char from element and move to anther element when you read the second char
# do it k times

while i < k:
    for element in clear_text_list:
        encrypted_text+=element[i]
    i+=1

with open("cipher_text.txt", "w") as f:
    f.write(encrypted_text)

print("\nFinally, text encrypted by Rail Fence Transposition Cipher is: ")
print(encrypted_text)