import sys
from hashlib import sha256, sha384
import os
import random

byteorder = sys.byteorder


def get_power_2_factors(n: int) -> (int, int):
    r = 0
    d = n
    while n > 0 and d % 2 == 0:
        d = d // 2
        r += 1
    return r, d


def miller_rabin_prime_test(n: int, k: int = 64) -> bool:
    # Factor powers of 2 from n - 1 s.t. n - 1 = 2^r * d
    r, d = get_power_2_factors(n - 1)

    for i in range(k):
        a = get_random_bits(n.bit_length())
        while a not in range(2, n - 2 + 1):
            a = get_random_bits(n.bit_length())
        x = pow(a, d, n)
        if x == 1 or x == n - 1:
            continue
        n_1_found = False
        for j in range(r - 1):
            x = pow(x, 2, n)
            if x == n - 1:
                n_1_found = True
                break
        if not n_1_found:
            return False
    return True


def get_random_bits(bit_length: int) -> int:
    return int.from_bytes(os.urandom((bit_length + 7) // 8), 'big')


def bezpieczny_randint(poczatek_inkluzywny, koniec_ekskluzywny):
    dlugosc_przedzialu = koniec_ekskluzywny - poczatek_inkluzywny
    maksymalna_liczba_bitow_zapisu = dlugosc_przedzialu.bit_length()
    # liczba miedzy 0 a 2**maksymalna_liczba_bitow_zapisu
    wyn = get_random_bits(maksymalna_liczba_bitow_zapisu)
    # liczba miedzy 0 a dlugosc_przedzialu
    while wyn >= dlugosc_przedzialu:
        wyn = get_random_bits(maksymalna_liczba_bitow_zapisu)
    # liczba miedzy poczatek_inkluzywny a koniec_ekskluzywny
    wyn = poczatek_inkluzywny + wyn
    return wyn


def generate_prime_number(bit_length: int) -> int:
    # prime needs to be in range [2^(n-1), 2^n-1]
    low = pow(2, bit_length - 1)
    high = pow(2, bit_length) - 1

    while True:

        # Generate odd prime candidate in range
        candidate_prime = get_random_bits(bit_length)
        while candidate_prime not in range(low, high + 1) or not candidate_prime % 2:
            candidate_prime = get_random_bits(bit_length)

        # with k rounds, miller rabin test gives false positive with probability (1/4)^k = 1/(2^2k)
        k = 64
        if miller_rabin_prime_test(candidate_prime, k):
            return candidate_prime


def rozszerzony_algorytm_euklidesa(a0, b0):
    a, b = a0, b0
    x, y = 1, 0
    r, s = 0, 1
    while b != 0:
        c = a % b
        q = a // b
        a, b = b, c
        r_pom, s_pom = r, s
        r = x - q * r
        s = y - q * s
        x, y = r_pom, s_pom
    nwd = x * a0 + y * b0
    return a, x, y, nwd


def NWD(a, b):
    nwd1, _, _, nwd2 = rozszerzony_algorytm_euklidesa(a, b)
    if nwd1 == nwd2:
        return nwd1
    else:
        return None


def odwrotnosc_modulo(a, n):
    _, x, y, _ = rozszerzony_algorytm_euklidesa(a, n)
    if a != 1 and x == 1:
        odwrotnosc_a = "nie ma elementu odwrotnego"
    elif x > 0:
        odwrotnosc_a = x % n
    else:
        odwrotnosc_a = (n + x) % n
    return odwrotnosc_a


def szybkie_potegowanie_modulo(podstawa, wykladnik, rzad):
    pom = 1
    while wykladnik > 1:
        ostatni_bit = wykladnik & 1
        wykladnik >>= 1
        if ostatni_bit == 1:
            pom = (pom * podstawa) % rzad
        podstawa = (podstawa * podstawa) % rzad
    return (podstawa * pom) % rzad


def genDSA():
    q = generate_prime_number(160)
    print("q", q.bit_length())
    t = bezpieczny_randint(0, 9)
    print("t", t, 512 + 64 * t)
    k = 2 ** (511 - 160 + 1 + 64 * t)
    while True:
        p = k * q + 1
        if miller_rabin_prime_test(p, 64):
            break
        k += 1
    print("p", p.bit_length())
    while True:
        g = bezpieczny_randint(2 ** (p.bit_length() - 1), 2 ** (p.bit_length()))
        alfa = szybkie_potegowanie_modulo(g, (p - 1) // q, p)
        if alfa != 1:
            break
    print("g, alfa", g.bit_length(), alfa.bit_length())
    a = bezpieczny_randint(1, q)
    y = szybkie_potegowanie_modulo(alfa, a, p)
    # tlumaczenie symboli
    p, q, g, y, x = p, q, alfa, y, a
    klucz_publiczny_DSA = (p, q, g, y)
    klucz_prywatny_DSA = x
    return klucz_publiczny_DSA, klucz_prywatny_DSA


def sig_DSA(p, q, g, x, m):
    byteorder = sys.byteorder
    k = get_random_bits(q.bit_length())
    print("k", k)
    r = szybkie_potegowanie_modulo(g, k, p) % q
    print("r", r)
    H_od_m = sha256(m).digest()
    int_H_od_m = int.from_bytes(H_od_m, byteorder)
    s = (odwrotnosc_modulo(k, q) * (int_H_od_m + x * r)) % q
    print("s", s)
    return r, s


def ver_DSA(p, q, g, y, m, r, s):
    byteorder = sys.byteorder
    if not (0 < r and s < q):
        print("zle")
    H_od_m = sha256(m).digest()
    int_H_od_m = int.from_bytes(H_od_m, byteorder)
    u1 = (odwrotnosc_modulo(s, q) * int_H_od_m) % q
    u2 = (r * odwrotnosc_modulo(s, q)) % q
    v = ((szybkie_potegowanie_modulo(g, u1, p) * szybkie_potegowanie_modulo(y, u2, p)) % p) % q
    print("r", r)
    print("v", v)
    if v == r:
        b = 1
    else:
        b = 0
    return b


print("DSA")
klucz_publiczny_DSA, klucz_prywatny_DSA = genDSA()
print(klucz_publiczny_DSA, klucz_prywatny_DSA)
p, q, g, y = klucz_publiczny_DSA
x = klucz_prywatny_DSA

m = "Thats my Kung Fu".encode("ASCII")
print(m)
r, s = sig_DSA(p, q, g, x, m)
print(r, s)
b = ver_DSA(p, q, g, y, m, r, s)
print(b)