Affine Cipher Solver

Affine_Cipher_Solver.py

import re
import gmpy2
from gmpy2 import mpz, divm, powmod, gcd, gcdext, invert
import itertools
import operator
from itertools import combinations, permutations, groupby

## solve system of equation:
# a * l1 + b = l1_p mod 26
# a * l2 + b = l2_p mod 26
def solve_affine_enc_key(l1,l2, l1_p, l2_p):
    l1, l2, l1_p, l2_p = ord(l1) - ord('a') , ord(l2) - ord('a'), ord(l1_p) - ord('a'), ord(l2_p) - ord('a')
    # a (l1-l2) = l1_p - l2_p mod 26
    x = invert(l1-l2, 26)
    a = (x * (l1_p - l2_p)) % 26
    b = (l1_p - a*l1) % 26
    return (int(a),int(b))

def affine_dec_key(a, b):
    c = invert(a, 26)
    d = (-b * c) % 26
    return (int(c), int(d))

def affine_transform(ciphertext, a, b):
    message = ""
    for c in  ciphertext.lower():
        x = ord(c) - ord('a')
        m = chr( ((a * x + b) % 26 )+ ord('a') )
        message += m
    return message



reference_hi_freq_letters=['e', 't', 'a', 'o', 'i', 'n', 's', 'r']

def solve_affine(ciphertext):
    clst = sorted([(k, len(list(g)) ) for k, g in groupby(sorted(list(ciphertext.lower()))) ], key=operator.itemgetter(1), reverse=True)
    print(f"clst={clst}")
    hi_freq_letters = [c for c, f in clst[:3] ]
    for l1_prim, l2_prim in permutations(hi_freq_letters, 2):
        for l1, l2 in combinations(reference_hi_freq_letters[:3], 2):
            try:
                a,b = solve_affine_enc_key(l1, l2, l1_prim, l2_prim)
                c, d = affine_dec_key(a, b)
                message = affine_transform(ciphertext, c, d)
                yield {'message':message, 'key':(a,b), 'x':f"{l1}->{l1_prim} {l2}->{l2_prim}" }
            except:
                continue


s = """
LJUNV  PCGIP  HUTCD  JU
PWD  UJGTP  GGXUT  HUTC
D  LJUNV  PCGIP  ZXLBP 
AWDTJ  JAUPR  RLJUN  VP
CGI  PRGXU  TDWDT  PUPC
J  RXGDH  UAGIP  QXHCW 
DRYPH BXUTG IPGCV GI
"""

s = re.sub(r"\s+", "", s)

for m in solve_affine(s):
    print(m)