#!/usr/bin/env python3
import sys
from hashlib import sha1
from random import randint


CRYPTOPALS_p = 0x800000000000000089E1855218A0E7DAC38136FFAFA72EDA7859F2171E25E65EAC698C1702578B07DC2A1076DA241C76C62D374D8389EA5AEFFD3226A0530CC565F3BF6B50929139EBEAC04F48C3C84AFB796D61E5A4F9A8FDA812AB59494232C7D2B4DEB50AA18EE9E132BFA85AC4374D7F9091ABC3D015EFC871A584471BB1
CRYPTOPALS_q = 0xF4F47F05794B256174BBA6E9B396A7707E563C5B
CRYPTOPALS_g = 0x5958C9D3898B224B12672C0B98E06C60DF923CB8BC999D119458FEF538B8FA4046C8DB53039DB620C094C9FA077EF389B5322A559946A71903F990F1F7E0E025E2D7F7CF494AFF1A0470F5B64C36B625A097F1651FE775323556FE00B3608C887892878480E99041BE601A62166CA6894BDD41A7054EC89F756BA9FC95302291


def random_bytes(size):
    i = randint(0, 2 ** (8 * size))
    return i.to_bytes(size, sys.byteorder)


def invmod(a, n):
    t = 0
    new_t = 1
    r = n
    new_r = a
    while new_r != 0:
        quotient = r // new_r
        t, new_t = new_t, t - quotient * new_t
        r, new_r = new_r, r - quotient * new_r
    if r > 1:  # a is not invertible
        return None
    if t < 0:
        t += n
    return t


def int_byte_size(i):
    return (i.bit_length() + 7) // 8


class dsa:
    def set_default_parms(self):
        self.p = CRYPTOPALS_p
        self.q = CRYPTOPALS_q
        self.g = CRYPTOPALS_g

    def generate_keypair(self, p=CRYPTOPALS_p, q=CRYPTOPALS_q, g=CRYPTOPALS_g):
        self.p = p
        self.q = q
        self.g = g
        self._private_key = randint(1, q - 1)
        self.public_key = pow(g, self._private_key, p)

    def import_public_key(self, n, p=CRYPTOPALS_p, q=CRYPTOPALS_q, g=CRYPTOPALS_g):
        self.p = p
        self.q = q
        self.g = g
        self.public_key = n

    def _sign(self, hash, k):
        r = pow(self.g, k, self.p) % self.q
        inverse = invmod(k, self.q)
        s = (inverse * (hash + (self._private_key * r))) % self.q
        return r, s

    def sign_hash(self, hash):
        r = 0
        s = 0
        while True:  # Need to retry on 0'd outputs
            k = randint(1, self.q - 1)
            r, s = self._sign(hash, k)
            # Looks like the safetys need to be disabled for testing g=0
            # if r == 0 or s == 0:
            if s == 0:
                continue  # Try a new value for k
            break
        return r, s

    def sign(self, data):
        return self.sign_hash(int.from_bytes(sha1(data).digest(), "big"))

    def verify(self, data, signature):
        r = signature[0]
        s = signature[1]
        # Looks like the safetys need to be disabled for testing g=0
        # if r <= 0 or r >= self.q:
        #     return False
        if s <= 0 or s >= self.q:
            return False

        hash = int.from_bytes(sha1(data).digest(), "big")
        w = invmod(s, self.q)
        u1 = (hash * w) % self.q
        u2 = (r * w) % self.q
        g_u1 = pow(self.g, u1, self.p)
        g_u2 = pow(self.public_key, u2, self.p)
        v = ((g_u1 * g_u2) % self.p) % self.q
        return v == r


def generate_magic_sign_for_p_plus_1(p, q, public_key, data):
    hash = int.from_bytes(sha1(data).digest(), "big")
    r = pow(public_key, hash, p) % q
    s = r * invmod(hash, q) % q
    return r, s


def test_g(g):
    _dsa = dsa()
    _dsa.generate_keypair(g=g)
    for _ in range(10):
        s = _dsa.sign(random_bytes(randint(10, 128)))
        if not _dsa.verify(random_bytes(randint(10, 128)), s):
            print("Unable to validate any signature against any data")
            return False
    return True


def test_p_plus_1():
    _dsa = dsa()
    _dsa.generate_keypair(g=CRYPTOPALS_p + 1)
    for _ in range(10):
        d = random_bytes(randint(10, 128))
        s = generate_magic_sign_for_p_plus_1(
            CRYPTOPALS_p, CRYPTOPALS_q, _dsa.public_key, d
        )
        if not _dsa.verify(d, s):
            print("Unable to validate magic signature")
            return False
    return True


for g in [0, 1]:
    if test_g(g):
        print(f"Sign/Verify with g={g} Passed!")
    else:
        print(f"Sign/Verify with g={g} Failed!")

if test_p_plus_1():
    print("Magic Signature with g=p+1 Passed!")
else:
    print("Magic Signature with g=p+1 Failed!")
