#!/usr/bin/python3
# -*- coding: utf-8 -*-

# Public Key Cipher
# https://www.nostarch.com/crackingcodes/ (BSD Licensed)

import sys, math
from typing import List, Optional, Tuple

# The public and private keys for this program are created by
# the makePublicPrivateKeys.py program.
# This program must be run in the same folder as the key files.

SYMBOLS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz1234567890 !?."


def main() -> None:
    # Runs a test that encrypts a message to a file or decrypts a message from a file.
    filename = "encrypted_file.txt"  # The file to write to/read from.
    # mode = "decrypt"  # Set to either 'encrypt' or 'decrypt'.
    mode = input("Type 'encrypt' or 'decrypt' and hit enter: ")

    if mode == "encrypt":
        message = "Journalists belong in the gutter because that is where the ruling classes throw their guilty secrets. Gerald Priestland. The Founding Fathers gave the free press the protection it must have to bare the secrets of government and inform the people. Hugo Black."
        pubKeyFilename = "al_sweigart_pubkey.txt"
        print("Encrypting and writing to %s..." % (filename))
        encryptedText = encryptAndWriteToFile(filename, pubKeyFilename, message)

        print("Encrypted text:")
        print(encryptedText)

    elif mode == "decrypt":
        privKeyFilename = "al_sweigart_privkey.txt"
        print("Reading from %s and decrypting..." % (filename))
        decryptedText = readFromFileAndDecrypt(filename, privKeyFilename)

        print("Decrypted text:")
        print(decryptedText)


def getBlocksFromText(message: str, blockSize: int) -> List[int]:
    # Converts a string message to a list of block integers.
    for character in message:
        if character not in SYMBOLS:
            print("ERROR: The symbol set does not have the character %s" % (character))
            sys.exit()
    blockInts = []
    for blockStart in range(0, len(message), blockSize):
        # Calculate the block integer for this block of text:
        blockInt = 0
        for i in range(blockStart, min(blockStart + blockSize, len(message))):
            blockInt += (SYMBOLS.index(message[i])) * (len(SYMBOLS) ** (i % blockSize))
        blockInts.append(blockInt)
    return blockInts


def getTextFromBlocks(blockInts: List[int], messageLength: int, blockSize: int) -> str:
    # Converts a list of block integers to the original message string.
    # The original message length is needed to properly convert the last
    # block integer.
    message: List[str] = []
    for blockInt in blockInts:
        blockMessage: List[str] = []
        for i in range(blockSize - 1, -1, -1):
            if len(message) + i < messageLength:
                # Decode the message string for the 128 (or whatever
                # blockSize is set to) characters from this block integer:
                charIndex = blockInt // (len(SYMBOLS) ** i)
                blockInt = blockInt % (len(SYMBOLS) ** i)
                blockMessage.insert(0, SYMBOLS[charIndex])
        message.extend(blockMessage)
    return "".join(message)


def encryptMessage(message: str, key: Tuple[int, int], blockSize: int) -> List[int]:
    # Converts the message string into a list of block integers, and then
    # encrypts each block integer. Pass the PUBLIC key to encrypt.
    encryptedBlocks = []
    n, e = key

    for block in getBlocksFromText(message, blockSize):
        # ciphertext = plaintext ^ e mod n
        encryptedBlocks.append(pow(block, e, n))
    return encryptedBlocks


def decryptMessage(
    encryptedBlocks: List[int], messageLength: int, key: Tuple[int, int], blockSize: int
) -> str:
    # Decrypts a list of encrypted block ints into the original message
    # string. The original message length is required to properly decrypt
    # the last block. Be sure to pass the PRIVATE key to decrypt.
    decryptedBlocks = []
    n, d = key
    for block in encryptedBlocks:
        # plaintext = ciphertext ^ d mod n
        decryptedBlocks.append(pow(block, d, n))
    return getTextFromBlocks(decryptedBlocks, messageLength, blockSize)


def readKeyFile(keyFilename: str) -> Tuple[int, int, int]:
    # Given the filename of a file that contains a public or private key,
    # return the key as a (n,e) or (n,d) tuple value.
    with open(keyFilename) as fo:
        content = fo.read()
    keySize, n, EorD = content.split(",")
    return (int(keySize), int(n), int(EorD))


def encryptAndWriteToFile(
    messageFilename: str,
    keyFilename: str,
    message: str,
    blockSize: Optional[int] = None,
) -> str:
    # Using a key from a key file, encrypt the message and save it to a
    # file. Returns the encrypted message string.
    keySize, n, e = readKeyFile(keyFilename)
    if blockSize is None:
        # If blockSize isn't given, set it to the largest size allowed by the key size and symbol set size.
        blockSize = int(math.log(2**keySize, len(SYMBOLS)))
    # Check that key size is large enough for the block size:
    if not (math.log(2**keySize, len(SYMBOLS)) >= blockSize):
        sys.exit(
            "ERROR: Block size is too large for the key and symbol set size. Did you specify the correct key file and encrypted file?"
        )
    # Encrypt the message:
    encryptedBlocks = encryptMessage(message, (n, e), blockSize)

    # Convert the large int values to one string value:
    # for i in range(len(encryptedBlocks)):
    #     encryptedBlocks[i] = str(encryptedBlocks[i])
    # encryptedContent = ",".join(encryptedBlocks)
    encryptedContent = ",".join(map(str, encryptedBlocks))

    # Write out the encrypted string to the output file:
    encryptedContent = "%s_%s_%s" % (len(message), blockSize, encryptedContent)
    with open(messageFilename, "w") as fo:
        fo.write(encryptedContent)
    # Also return the encrypted string:
    return encryptedContent


def readFromFileAndDecrypt(messageFilename: str, keyFilename: str) -> str:
    # Using a key from a key file, read an encrypted message from a file
    # and then decrypt it. Returns the decrypted message string.
    keySize, n, d = readKeyFile(keyFilename)

    # Read in the message length and the encrypted message from the file:
    with open(messageFilename) as fo:
        content = fo.read()
    messageLength_str, blockSize_str, encryptedMessage = content.split("_")
    messageLength = int(messageLength_str)
    blockSize = int(blockSize_str)

    # Check that key size is large enough for the block size:
    if not math.log(2**keySize, len(SYMBOLS)) >= blockSize:
        sys.exit(
            "ERROR: Block size is too large for the key and symbol set size. Did you specify the correct key file and encrypted file?"
        )

    # Convert the encrypted message into large int values:
    encryptedBlocks = []
    for block in encryptedMessage.split(","):
        encryptedBlocks.append(int(block))

    # Decrypt the large int values:
    return decryptMessage(encryptedBlocks, messageLength, (n, d), blockSize)


# If publicKeyCipher.py is run (instead of imported as a module) call
# the main() function.
if __name__ == "__main__":
    main()