affine_cipher.py

importrandom
importsys

frommaths.greatest_common_divisorimportgcd_by_iterative

from . importcryptomath_moduleascryptomath

SYMBOLS= (
r""" !"#$%&'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\]^_`"""
r"""abcdefghijklmnopqrstuvwxyz{|}~"""
)


defcheck_keys(key_a: int, key_b: int, mode: str) ->None:
ifmode=="encrypt":
ifkey_a==1:
sys.exit(
"The affine cipher becomes weak when key "
"A is set to 1. Choose different key"
 )
ifkey_b==0:
sys.exit(
"The affine cipher becomes weak when key "
"B is set to 0. Choose different key"
 )
ifkey_a<0orkey_b<0orkey_b>len(SYMBOLS) -1:
sys.exit(
"Key A must be greater than 0 and key B must "
f"be between 0 and {len(SYMBOLS) -1}."
 )
ifgcd_by_iterative(key_a, len(SYMBOLS)) !=1:
sys.exit(
f"Key A {key_a} and the symbol set size {len(SYMBOLS)} "
"are not relatively prime. Choose a different key."
 )


defencrypt_message(key: int, message: str) ->str:
"""
 >>> encrypt_message(4545, 'The affine cipher is a type of monoalphabetic '
 ... 'substitution cipher.')
 'VL}p MM{I}p~{HL}Gp{vp pFsH}pxMpyxIx JHL O}F{~pvuOvF{FuF{xIp~{HL}Gi'
 """
key_a, key_b=divmod(key, len(SYMBOLS))
check_keys(key_a, key_b, "encrypt")
cipher_text=""
forsymbolinmessage:
ifsymbolinSYMBOLS:
sym_index=SYMBOLS.find(symbol)
cipher_text+=SYMBOLS[(sym_index*key_a+key_b) %len(SYMBOLS)]
else:
cipher_text+=symbol
returncipher_text


defdecrypt_message(key: int, message: str) ->str:
"""
 >>> decrypt_message(4545, 'VL}p MM{I}p~{HL}Gp{vp pFsH}pxMpyxIx JHL O}F{~pvuOvF{FuF'
 ... '{xIp~{HL}Gi')
 'The affine cipher is a type of monoalphabetic substitution cipher.'
 """
key_a, key_b=divmod(key, len(SYMBOLS))
check_keys(key_a, key_b, "decrypt")
plain_text=""
mod_inverse_of_key_a=cryptomath.find_mod_inverse(key_a, len(SYMBOLS))
forsymbolinmessage:
ifsymbolinSYMBOLS:
sym_index=SYMBOLS.find(symbol)
plain_text+=SYMBOLS[
 (sym_index-key_b) *mod_inverse_of_key_a%len(SYMBOLS)
 ]
else:
plain_text+=symbol
returnplain_text


defget_random_key() ->int:
whileTrue:
key_b=random.randint(2, len(SYMBOLS))
key_b=random.randint(2, len(SYMBOLS))
ifgcd_by_iterative(key_b, len(SYMBOLS)) ==1andkey_b%len(SYMBOLS) !=0:
returnkey_b*len(SYMBOLS) +key_b


defmain() ->None:
"""
 >>> key = get_random_key()
 >>> msg = "This is a test!"
 >>> decrypt_message(key, encrypt_message(key, msg)) == msg
 True
 """
message=input("Enter message: ").strip()
key=int(input("Enter key [2000 - 9000]: ").strip())
mode=input("Encrypt/Decrypt [E/D]: ").strip().lower()

ifmode.startswith("e"):
mode="encrypt"
translated=encrypt_message(key, message)
elifmode.startswith("d"):
mode="decrypt"
translated=decrypt_message(key, message)
print(f"\n{mode.title()}ed text: \n{translated}")


if__name__=="__main__":
importdoctest

doctest.testmod()
# main()