# Rozszerzony algorytm Euklidesa do znajdowania odwrotności modulo
# Przykład użycia: modinv(7, 26) = 15 
def egcd(a, b): 
    x,y, u,v = 0,1, 1,0
    while a != 0: 
        q, r = b//a, b%a 
        m, n = x-u*q, y-v*q 
        b,a, x,y, u,v = a,r, u,v, m,n 
    gcd = b 
    return gcd, x, y 
  
def modinv(a, m): 
    gcd, x, y = egcd(a, m) 
    if gcd != 1: 
        return None # Odwrotność modulo nie istnieje 
    else: 
        return x % m 
  
  
# Funkcja szyfrująca, szyfr afiniczny 
# Zwraca szyfrogram
def encrypt(text, key): 
    ''' 
    E = (a * x + b) % 26 
    '''
    return ''.join([ chr((( key[0]*(ord(t) - ord('A')) + key[1] ) % 26)  
                  + ord('A')) for t in text.upper().replace(' ', '') ]) 
  
  
# Funkcja odszyfrowująca, szyfr afiniczny
# Zwraca pierwotną wiadomość
def decrypt(cipher, key): 
    ''' 
    D = (a^-1 * (x - b)) % 26 
    '''
    return ''.join([ chr((( modinv(key[0], 26)*(ord(c) - ord('A') - key[1]))  
                    % 26) + ord('A')) for c in cipher ]) 
  
  
# Test funkcji szyfrującej i odszyfrowującej
def main(): 
    # Deklaracja tekstu i klucza
    text = 'CODEBOOK'
    key = [17, 20] 
  
    # Wywołanie funkcji szyfrującej
    encrypted_text = encrypt(text, key) 
  
    print('Zaszyfrowana wiadomość: {}'.format(encrypted_text )) 
  
    # Wywołanie funkcji odszyfrowującej
    print('Odszyfrowana wiadomość: {}'.format
    (decrypt(encrypted_text, key) )) 
  
  
if __name__ == '__main__': 
    main() 
