#!/usr/bin/python

from __future__ import absolute_import, division, print_function

import subprocess
from builtins import bytes, range
from os.path import abspath, dirname
from os.path import join as join_path
from random import randint


from CryptoAttacks.Block.gcm import *
from CryptoAttacks.Utils import log


def test_polynomials():
    print("Test polynomials")
    Pmod = GF_2k_generator(128, [128,7,2,1,0])
    P = Pmod(0b10011010101100110100100110011101100110010111111000111011101000000110110100010101000101100100111100011001010100100110100111011000)
    Q = Pmod(0b01111010101010110111000011011100010011101111000001010000011000010000111010001111100001111010110001001000011101000011111110010101)
    print(P.to_bits(), bin(P.to_int()), P)
    print(Q.to_bits(), bin(Q.to_int()), Q)
    w = P*Q
    print(w.to_bits(), bin(w.to_int()), w)
    assert Q.coefficients == Pmod(Q.coefficients).coefficients
    assert Q.coefficients == Pmod(Q.to_int()).coefficients
    assert Q.coefficients == Pmod(Q.to_bytes()).coefficients
    print('')

    Pmod = GF_2k_generator(10, [11,7,2,1,0])
    c1 = Pmod(1)
    c2 = Pmod(0)
    c3 = Pmod(0)
    c4 = Pmod(0)
    polynomial1 = Polynomial_128([c1,c2,c3,c4])

    c1 = Pmod(1236)
    c2 = Pmod(0)
    c3 = Pmod(0)
    c4 = Pmod(0)
    polynomial2 = Polynomial_128([c1,c2,c3,c4])
    print(polynomial1)
    print(polynomial2)
    print("+", polynomial1 + polynomial2)
    print("*", polynomial1 * polynomial2)
    q = polynomial1 / polynomial2
    r = polynomial1 % polynomial2
    print("/", q)
    print("%", r)
    print('')
    print(polynomial1)
    print(polynomial2*q + r)
    print('')


def test_gcm():
    print("Test GCM")
    plaintext = bytes(b'hn9YA(F BW&B (W&&W(RT&WEF f7*WB FTgsdc')
    additional = bytes(b'j gej8g0SRYH8s  8s9yf sgd78taDS* GASyd ')
    key = bytes(b'xgrtjdh&LA28XNwh')
    nonce = bytes(b'a drO*1@((js')
    ciphertext, tag = gcm_encrypt(plaintext, additional, key, nonce)
    assert gcm_verify(tag, ciphertext, additional, key, nonce)

    blocks = aes_bytes_to_poly_blocks(ciphertext, additional)
    ciphertext2, additional2 = poly_blocks_to_aes_bytes(blocks)
    assert ciphertext == ciphertext2
    assert additional == additional2


def polynomial_factors_product(factorization):
    """factorization: [(poly1, power), (poly2, power)]"""
    result = factorization[0][0].one_element()
    for f, f_degree in factorization:
      result *= f**f_degree
    return result


def test_factor():
    print("Test factor")
    Pmod = GF_2k_generator(9, [9,7,2,1,0])
    c1 = Pmod(31)
    c2 = Pmod(0)
    c3 = Pmod(0)
    c4 = Pmod(3)
    polynomial1 = Polynomial_128([c1,c2,c3,c4])

    c1 = Pmod(237)
    c2 = Pmod(1)
    c3 = Pmod(0)
    c4 = Pmod(10)
    polynomial2 = Polynomial_128([c1,c2,c3,c4])
    polynomial = polynomial1 * polynomial2
    print(polynomial1)
    print(polynomial2)
    print(polynomial)
    print(polynomial.monic())
    print('')

    factorization = factor_polynomial(polynomial)
    print(factorization)
    result = polynomial.one_element()
    for f, f_degree in factorization:
      result *= f**f_degree
    print(result)
    print('')

    assert polynomial_factors_product(factorization) == polynomial.monic()


def test_repeated_nonce():
    print("Test Key-Recovery Attack on GCM with Repeated Nonces")
    for _ in range(3):
        nonce = random_bytes(12)
        key = random_bytes(16)

        h = bytes(AES.new(key, AES.MODE_ECB).encrypt(bytes(b'\x00'*16)))
        h = aes_polynomial(h)
        
        ciphertexts_additionals_tags = []
        for _ in range(4):
            plaintext = random_bytes(randint(0, 50))
            additional = random_bytes(randint(0, 50))
            ciphertext, tag = gcm_encrypt(plaintext, additional, key, nonce)
            ciphertexts_additionals_tags.append((ciphertext, additional, tag))

        valid_ciphertext, valid_additional, valid_tag = ciphertexts_additionals_tags[0]
        auth_key_candidates = recover_key_repated_nonce(ciphertexts_additionals_tags)
        assert h.to_bytes() in auth_key_candidates

        # try found auth key candidates
        correct_auth_key_found = False
        for auth_key in auth_key_candidates:
            forged_ciphertext = random_bytes(randint(0, 10))
            forged_additional = random_bytes(randint(0, 10))
            forged_tag = gcm_forge_tag(ciphertext=forged_ciphertext, additional=forged_additional, auth_key=auth_key,
                valid_ciphertext=valid_ciphertext, valid_additional=valid_additional, valid_tag=valid_tag)

            if gcm_verify(forged_tag, forged_ciphertext, forged_additional, key, nonce):
                correct_auth_key_found = True
                break

        assert correct_auth_key_found


def run():
    log.level = 'debug'

    test_polynomials()
    test_gcm()
    test_factor()
    test_repeated_nonce()


if __name__ == "__main__":
    run()

[((58, 16, 58, 64), 'builtins.bytes', 'bytes', ({(58, 22, 58, 63): "b'hn9YA(F BW&B (W&&W(RT&WEF f7*WB FTgsdc'"}, {}), "(b'hn9YA(F BW&B (W&&W(RT&WEF f7*WB FTgsdc')", False, 'from builtins import bytes, range\n'), ((59, 17, 59, 66), 'builtins.bytes', 'bytes', ({(59, 23, 59, 65): "b'j gej8g0SRYH8s  8s9yf sgd78taDS* GASyd '"}, {}), "(b'j gej8g0SRYH8s  8s9yf sgd78taDS* GASyd ')", False, 'from builtins import bytes, range\n'), ((60, 10, 60, 36), 'builtins.bytes', 'bytes', ({(60, 16, 60, 35): "b'xgrtjdh&LA28XNwh'"}, {}), "(b'xgrtjdh&LA28XNwh')", False, 'from builtins import bytes, range\n'), ((61, 12, 61, 34), 'builtins.bytes', 'bytes', ({(61, 18, 61, 33): "b'a drO*1@((js'"}, {}), "(b'a drO*1@((js')", False, 'from builtins import bytes, range\n'), ((113, 13, 113, 21), 'builtins.range', 'range', ({(113, 19, 113, 20): '(3)'}, {}), '(3)', False, 'from builtins import bytes, range\n'), ((121, 17, 121, 25), 'builtins.range', 'range', ({(121, 23, 121, 24): '(4)'}, {}), '(4)', False, 'from builtins import bytes, range\n'), ((117, 53, 117, 70), 'builtins.bytes', 'bytes', ({(117, 59, 117, 69): "b'\\x00' * 16"}, {}), "(b'\\x00' * 16)", False, 'from builtins import bytes, range\n'), ((122, 37, 122, 51), 'random.randint', 'randint', ({(122, 45, 122, 46): '0', (122, 48, 122, 50): '50'}, {}), '(0, 50)', False, 'from random import randint\n'), ((123, 38, 123, 52), 'random.randint', 'randint', ({(123, 46, 123, 47): '0', (123, 49, 123, 51): '50'}, {}), '(0, 50)', False, 'from random import randint\n'), ((134, 45, 134, 59), 'random.randint', 'randint', ({(134, 53, 134, 54): '0', (134, 56, 134, 58): '10'}, {}), '(0, 10)', False, 'from random import randint\n'), ((135, 45, 135, 59), 'random.randint', 'randint', ({(135, 53, 135, 54): '0', (135, 56, 135, 58): '10'}, {}), '(0, 10)', False, 'from random import randint\n')]