#@# Digital signatures
#@# ==================

#@# Public-key cryptography (= asymmatric cryptography) has various goals:
#@#    1. encryption (e.g. RSA, see: 6rsa.sage; ElGamal, see: 7diffie.sage);
#@#    2. key exchange (e.g. Diffie-Hellman key exchange, see: 7diffie.sage);
#@#    3. digital signature:
#@#       a cryptographic method to "sign" digital messages or documents,
#@#       so that later it can be verified by other parties.
#@#
#@# Digital signatures consist of 3 algorithms:
#@# 1. Gen() -> (pk, sk): key generation, which returns:
#@#    * a public key (pk) and
#@#    * a private key (sk).
#@# 2. Sign(m, sk) -> s: signing algorithm, which returns a signature ('s'),
#@#    a value calculated from the message using the _private_ key.
#@# 3. Verify(m, s, pk): verification algorithm, which returns True/False,
#@#    depending on whether the signature is valid for that message,
#@#    i.e. it was indeed signed by the holder of the private key. Note:
#@#    verification only requires the _public_ key, i.e. anyone can do it.
#@#
#@# Common digital signature algorithms:
#@#    * RSA (see exercises 1-3);
#@#    * DSA = Digital Signature Algorithm (see exercises 4-5);
#@#    * ECDSA (variant of DSA with elliptic curves instead of exponentiation).
#@#
#@# Common applications of digital signatures:
#@#    * signing legal contracts (instead of signing it on paper by ink);
#@#    * signing software updates: to verify that it indeed came from the
#@#      manufacturer, not from a malicious source;
#@#    * TLS (= the S in HTTPS): to verify that we communicate with the correct
#@#      website (e.g. a bank's website vs someone's malicious copy of it).

# Helper functions: convert strings to/from list of bytes:
def str2bytes (str): # "ABCD" -> [65,66,67,68]
    return [ord(c) for c in str]
def bytes2str (bytes): # [65,66,67,68] -> "ABCD"
    return "".join([chr(b) for b in bytes])
# Helper functions: convert list of bytes to/from big integer (in base 256)
def bytes2int (b): # [11,22,33] -> 11*256^2 + 22*256 + 33 = 726561
    return ZZ(b[::-1], 256)
def int2bytes (i, length=0):
    return ZZ(i).digits(256, padto=length)[::-1]


#@# ---------------------------------------------------------------------------
#@# Exercise 1: plain RSA signature
#@#
#@# Plain RSA signature is the reverse of the plain RSA encryption.
#@# Key generation: same as with encryption, see 6rsa.sage.
#@# Signing (with the private key sk = (N,d)): s = m^d mod N.
#@# Verification (with the public key pk = (N,e)): check if s^e mod N == m.
#@#
#@# Implement plain RSA signature and verification as described above.
#@# (The key generation part is already copied from 6rsa.sage.)

# Generate an n-bit prime number (i.e. 2^(n-1) < p < 2^n).
def gen_prime (n):
    #return random_prime(2^n, lbound=2^(n-1))  # works, but not secure RNG
    import secrets
    while True:
        p = 2^(n-1) + secrets.randbelow(int(2^(n-1)))
        p |= 1 # make the last bit 1 (so that p is odd)
        if is_prime(p): # is_pseudoprime(p) is faster, but not always correct
            return p
# Generate RSA keys (with n-bit primes, i.e. 2n-bit RSA)
def rsa_gen (n):
    p = gen_prime(n)
    q = gen_prime(n)
    N = p*q
    phi_N = (p-1)*(q-1) # don't use euler_phi(N), it would try to factor N
    for e in range(3, phi_N, 2): # often e is fixed, and N is chosen for e
        if gcd(e, phi_N) == 1:
            break
    d = inverse_mod(e, phi_N)
    pk = (N,e)  # public key
    sk = (N,d)  # private key
    return (pk, sk)

# Signature
# m: the message to be signed (list of bytes)
# sk: the private key (N,d)
# Return: the signature s (integer)
def rsa_sign (m, sk):
    pass #TODO

# Verification
# m: the signed message (list of bytes)
# s: the signature returned by rsa_sign()
# pk: the public key (N,e)
# Return: is it valid? (True/False)
def rsa_verify (m, s, pk):
    pass #TODO


#@# Test any signature algorithm (used in many exercises)
def test_signature (pk, sk, m, sign, verify): #@#
    print("m =", repr(bytes2str(m))) #@#
    s = sign(m, sk) #@#
    print("s =", s) #@#
    ok = verify(m, s, pk) #@#
    print("verify:        ", ok, "!!" if not ok else "") #@#
    bad_s = s+1 if s in ZZ else (si+1 for si in s) #@#
    print("verify (bad s):", verify(m, bad_s, pk)) #@#
    m[0] += 1 #@#
    print("verify (bad m):", verify(m, s, pk)) #@#

#@# Test plain RSA signaure
def test_rsa (): #@#
    pk, sk = rsa_gen(128) #@# 256-bit RSA
    print("N =", pk[0]) #@#
    m = str2bytes("Message for signing (<32 bytes)") #@# 256/8 = 32
    test_signature(pk, sk, m, rsa_sign, rsa_verify) #@#
#test_rsa()

#@# Note: plain RSA signature has two problems:
#@# 1. it is not secure (see next exercise);
#@# 2. and it cannot sign messages longer than N.
#@# See exercise 3 for a possible solution of both problems.


#@# ---------------------------------------------------------------------------
#@# Exercise 2: forging plain RSA signature
#@#
#@# A signature scheme is broken if someone can "forge" a signature, i.e.
#@# construct a valid signature for a message without the private key.
#@# Plain RSA signature (like plain RSA encryption, see 6rsa.sage) is insecure.
#@# One way to forge signatures is to work backwards: pick a random s, and then
#@# calculate m = s^e mod N. This s, by definition, is a valid signature for m.
#@# This doesn't seem too useful though, because the attacker didn't choose m,
#@# so it is just a random-looking string that was calculated from a random s.
#@# However, a secure signature scheme shouldn't allow forging _any_ signature.
#@# Also, the attacker can repeatedly generate s until the calculated message m
#@# is "good enough", for whatever purpose the attacker has.
#@# To demonstrate this, assume a very short RSA with only 4 bytes of message,
#@# and find a signature for a message whose last two bytes are both 111.
#@# (Warning: int2bytes() removes leading zeros, i.e. m may be shorter than 4!)

# pk: the public key, (N, e)
# Return: (m, s), a message (list of bytes) with a valid signature (integer).
def forge_rsa (pk):
    pass #TODO
    #return m, s

def test_forge_rsa (): #@#
    pk, sk = rsa_gen(16) #@# 32-bit (= 4-byte) RSA (very short)
    print("N =", pk[0]) #@#
    m, s = forge_rsa(pk) #@#
    print("m =", m) #@# (should end with [..., 111, 111])
    print("s =", s) #@#
    print("verify:", rsa_verify(m, s, pk)) #@#
#test_forge_rsa()


#@# ---------------------------------------------------------------------------
#@# Exercise 3: hash-and-sign paradigm
#@#
#@# To fix the above problems with RSA signature, we can use a hash function
#@# (see 5hash.sage). Then, instead of calculating with the original message m,
#@# use H(m), the hash value of m. For example, the signature is calculated as:
#@#    s = H(m)^d mod N.
#@# This prevents the above attack, because the attacker would need to invert
#@# the hash function, which is not possible for a cryptographic hash function.
#@# Implement hashed RSA signature, using the below hash function.

#@# SHA-256 hash of a message (list of bytes), converted to an integer:
def int_hash (m): #@#
    import hashlib #@#
    return bytes2int(list(hashlib.sha256(bytes(m)).digest())) #@#

# Signature
def rsa_hash_sign (m, sk):
    pass #TODO

# Verification
def rsa_hash_verify (m, s, pk):
    pass #TODO

def test_rsa_hash (): #@#
    pk, sk = rsa_gen(129) #@# >256-bit RSA
    print("N =", pk[0]) #@#
    m = str2bytes("Message to be signed - now it can be as long as we want!") #@#
    test_signature(pk, sk, m, rsa_hash_sign, rsa_hash_verify) #@#
#test_rsa_hash()


#@# ---------------------------------------------------------------------------
#@# Exercise 4: Digital Signature Algoritm (DSA)
#@#
#@# Another signature scheme is DSA, which is similar to ElGamal encryption
#@# (see 7diffie.sage), both based on the Diffie-Hellman key exchange.
#@# But now, we need another prime number modulus, q, which divides p-1, where
#@# p is the original prime modulus. (Note: q is the modulus for exponents, as
#@# Fermat's little theorem allows changing the exponent mod p-1, not mod p.)
#@# (Note 2: in practice, ElGamal and Diffie-Hellman also often use this q.)
#@#
#@# DSA key generation (similar to ElGamal's Gen(), but now including q):
#@#    * Fixed constants: prime p, prime q | p-1, and base g of order q
#@#      (i.e. g^q mod p = 1). (Note: a code that finds these values is given.)
#@#    * Generate a random x between 1 and q-1.
#@#    * Calculate y := g^x mod p
#@#    * Public key:  pk = (p, q, g, y)
#@#    * Private key: sk = (p, q, g, x)
#@# DSA signature (s = Sign(m, sk)):
#@#    * Generate a random k between 1 and q-1.
#@#    * s1 := g^k mod p.
#@#    * s2 := k^(-1) * (H(m) + x*s1) mod q.
#@#        H(m): the hash of m (as an integer).
#@#        k^(-1): modular inverse of k mod q (see Sage's inverse_mod()).
#@#    * Return s = (s1, s2) as the signature.
#@# DSA verification (Verify(m, s, pk), where s = (s1, s2)):
#@#    * u1 := H(m) * s2^(-1) mod q
#@#    * u2 :=   s1 * s2^(-1) mod q
#@#    * v := g^u1 * y^u2 mod p
#@#    * Verification: is v == s1?
#@# Proof of correctness:
#@#    v = g^(H(m) * s2^(-1)) * (g^x)^(s1 * s2^(-1) =
#@#      = g^((H(m) + x*s1) * s2^(-1)) =
#@#      = g^((H(m) + x*s1) * (H(m) + x*s1)^-1 * k) = g^k = s1   (mod p)
#@#
#@# Implement Gen(), Sign() and Verify() for DSA, using int_hash() from above.

# Generate suitable p, q and g for DSA
def gen_dsa_params (n):
    import secrets
    g = 2
    while True:
        p = 2^(n-1) + secrets.randbelow(int(2^(n-1)))
        p |= 3 # make the last two bits 1 (so that p and q are odd)
        q = (p - 1) // 2  # p = 2*q + 1 where q is also prime
        if is_prime(p) and is_prime(q) and multiplicative_order(mod(g, p)) == q:
            return p, q, g

# Key generation
# Return: the pair (pk, sk) as described above.
def dsa_gen ():
    p, q, g = gen_dsa_params(128)
    pass #TODO

# Signature
# m: the plaintext (list of bytes)
# Return: s = (s1,s2), the signature (pair of integers)
def dsa_sign (m, sk):
    pass #TODO

# Verification
# Return: True/False
def dsa_verify (m, s, pk):
    pass #TODO

def test_dsa (): #@#
    pk, sk = dsa_gen() #@#
    m = str2bytes("A very important contract to be signed digitally") #@#
    test_signature(pk, sk, m, dsa_sign, dsa_verify) #@#
#test_dsa()


#@# ---------------------------------------------------------------------------
#@# Exercise 5: misusing k in DSA
#@#
#@# Random numbers can have various requirements in cryptography:
#@#    * unique: it is never reused (e.g. the IV in CTR mode, see 4modes.sage);
#@#    * unpredictable: it can't be guessed _beforehand_ (e.g. IV in CBC mode);
#@#    * secret: nobody but the owner(s) can know it (e.g. a private key).
#@# DSA's random number k has the strictest possible requirements: it must
#@# satisfy ALL of the above: it must be unique AND unpredictable AND secret.
#@# Violating any of these requirementes is fatal: it allows an attacker to
#@# recover the private key x, and so forge signatures for any message.

#@# a) Assume that k was leaked (i.e. the "secret" requirement was violated).
#@# Reminder: k was used in DSA's Sign() algorithm as follows:
#@#    s2 := k^(-1) * (H(m) + x*s1) mod q.
#@# Using the known k, solve this congruence for x, the private key.
#@# Note that the message (m) and the signature (s1, s2) are not secret.

# pk, m, s: the publicly known values (public key, message and signature).
# k: the leaked random number.
# Return: x, the private key.
def break_dsa_known_k (pk, m, s, k):
    pass #TODO

def test_break_dsa_1 (): #@#
    pk = (13383823505783419439, 6691911752891709719, 2, 8523134409469066290) #@#
    m = str2bytes("This message was signed by DSA - is it secure?") #@#
    rs = (5974756495726202401, 1066852202230446380) #@#
    assert dsa_verify(m, rs, pk) #@#
    k = 2361311613004974411 #@# <- leaked!
    print("broken x:", break_dsa_known_k(pk, m, rs, k)) #@#
#test_break_dsa_1()

#@# b) Assume that k was reused (i.e. the "unique" requirement was violated),
#@# i.e. an attacker learns two signed messages that used the same (unknown) k:
#@#     (s1, s2) = Sign(m1, sk),
#@#     (s1, s3) = Sign(m2, sk).
#@# (Note: reuse of k is obvious from the repeated s1, because s1 = g^k mod p.)
#@# The calculation of s2 (see above) and s3 (similarly) give two congruences,
#@# which can be solved for k, and this k can be used as above to solve for x.
#@# Implement this attack.

# pk: public key.
# m1, s1, s2: first message and its signature.
# m2, s1, s3: second message and its signature.
# Return: x, the private key.
def break_dsa_reused_k (pk, m1, m2, s1, s2, s3):
    pass #TODO

def test_break_dsa_2 (): #@#
    pk = (13383823505783419439, 6691911752891709719, 2, 10016235895338920060) #@#
    sk = (13383823505783419439, 6691911752891709719, 2, 1234567890123456789) #@#
    m1 = str2bytes("This is the first message signed by DSA") #@#
    m2 = str2bytes("This is the second message, reusing k") #@#
    s1,s2 = (265613433957746694, 5735286140954322700) #@#
    s1,s3 = (265613433957746694, 6504086051745092855) #@#
    assert dsa_verify(m1, (s1,s2), pk) #@#
    assert dsa_verify(m2, (s1,s3), pk) #@#
    print("broken x:", break_dsa_reused_k(pk, m1, m2, s1, s2, s3)) #@#
#test_break_dsa_2()

#@# Note: even a less serious misuse of k can fully compromise the private key.
#@# For example, if the random number generator is weak, and k is selected from
#@# only 2^15 = 32768 possibilities, then an attacker can just try all possible
#@# k values and do the above calculation. And this happened in real life:
#@# https://rdist.root.org/2009/05/17/the-debian-pgp-disaster-that-almost-was/