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

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

import sys, cryptomath, random
from typing import Tuple

SYMBOLS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz1234567890 !?."


def main() -> None:
    myMessage = """"A computer would deserve to be called intelligent \
if it could deceive a human into believing that it was human." -Alan Turing"""
    myKey = 2894
    myMode = "encrypt"  # Set to either 'encrypt' or 'decrypt'.

    if myMode == "encrypt":
        translated = encryptMessage(myKey, myMessage)
    elif myMode == "decrypt":
        translated = decryptMessage(myKey, myMessage)
    print("Key: %s" % (myKey))
    print("%sed text:" % (myMode.title()))
    print(translated)
    print("Full %sed text copied to clipboard." % (myMode))


def getKeyParts(key: int) -> Tuple[int, int]:
    """
    Splits a single integer key into two integers for Key A and Key B.
    The key to split is passed to the key parameter.
    Key A is calculated by using integer division,
    to divide key by len(SYMBOLS), the size of the symbol set.
    Integer division (//) returns the quotient without a remainder.
    The mod operator (%) on line 26 calculates the remainder,
    which we’ll use for Key B.
    For example,
    with 2894 as the key parameter and a SYMBOLS string of 66 characters,
    Key A would be 2894 // 66 = 43 and Key B would be 2894 % 66 = 56.
    To combine Key A and Key B back into a single key,
    multiply Key A by the size of the symbol set and add Key B to the product:
    (43 * 66) + 56 evaluates to 2894, which is the integer key we started with.
    This is not required, and is just a cute way to store 2 keys in one.
    """
    keyA = key // len(SYMBOLS)
    keyB = key % len(SYMBOLS)
    return (keyA, keyB)


def checkKeys(keyA: int, keyB: int, mode: str) -> None:
    if keyA == 1 and mode == "encrypt":
        sys.exit("Cipher is weak if key A is 1. Choose a different key.")
    if keyB == 0 and mode == "encrypt":
        sys.exit("Cipher is weak if key B is 0. Choose a different key.")
    if keyA < 0 or keyB < 0 or keyB > len(SYMBOLS) - 1:
        sys.exit(
            "Key A must be greater than 0 and Key B must be between 0 and %s."
            % (len(SYMBOLS) - 1)
        )
    if cryptomath.gcd(keyA, len(SYMBOLS)) != 1:
        sys.exit(
            "Key A (%s) and the symbol set size (%s) are not relatively prime."
            "Choose a different key." % (keyA, len(SYMBOLS))
        )


def encryptMessage(key: int, message: str) -> str:
    keyA, keyB = getKeyParts(key)
    checkKeys(keyA, keyB, "encrypt")
    ciphertext = ""
    for symbol in message:
        if symbol in SYMBOLS:
            # Encrypt the symbol:
            symbolIndex = SYMBOLS.find(symbol)
            ciphertext += SYMBOLS[(symbolIndex * keyA + keyB) % len(SYMBOLS)]
        else:
            ciphertext += symbol  # Append the symbol without encrypting.
    return ciphertext


def decryptMessage(key: int, message: str) -> str:
    keyA, keyB = getKeyParts(key)
    checkKeys(keyA, keyB, "decrypt")
    plaintext = ""
    modInverseOfKeyA = cryptomath.findModInverse(keyA, len(SYMBOLS))

    # just for sane typing on the findModInverse Optional[int]
    if modInverseOfKeyA:
        for symbol in message:
            if symbol in SYMBOLS:
                # Decrypt the symbol:
                symbolIndex = SYMBOLS.find(symbol)
                plaintext += SYMBOLS[
                    (symbolIndex - keyB) * modInverseOfKeyA % len(SYMBOLS)
                ]
            else:
                plaintext += symbol  # Append the symbol without decrypting.
    return plaintext


def getRandomKey() -> int:
    while True:
        keyA = random.randint(2, len(SYMBOLS))
        keyB = random.randint(2, len(SYMBOLS))
        if cryptomath.gcd(keyA, len(SYMBOLS)) == 1:
            return keyA * len(SYMBOLS) + keyB


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