import binascii

from flask import Flask, render_template, request, redirect, url_for
from flask_wtf import FlaskForm
from wtforms import StringField
from wtforms.validators import DataRequired, Length, ValidationError

from aesgcmanalysis import xor, gmac, gcm_encrypt, nonce_reuse_recover_secrets, gf128_to_bytes

app = Flask(__name__)

@app.route('/')
def index():
    return render_template('index.html')

def hex_check(form, field):
    if len(field.data) % 2 != 0:
        raise ValidationError(f'not valid hex; must have even length')
    if not all(c in '1234567890abcdef' for c in field.data):
        raise ValidationError(f'not valid hex; contains non-hex character')

def not_equal_to(other):
    def helper(form, field):
        print(other, form['m1'], field)
        if other not in form:
            return
        if form[other].data == field.data:
            raise ValidationError(f'must not be equal to {other}')
    return helper

class NonceReuseForm(FlaskForm):
    key = StringField('key', validators=[DataRequired(), Length(min=32, max=32), hex_check])
    nonce = StringField('nonce', validators=[DataRequired(), Length(min=24, max=24), hex_check])
    m1 = StringField('first message', validators=[DataRequired(), Length(min=1, max=64)])
    m2 = StringField('second message', validators=[DataRequired(), Length(min=1, max=64), not_equal_to('m1')])
    mf = StringField('forged message', validators=[DataRequired(), Length(min=1, max=64)])

@app.route('/nonce-reuse', methods=['GET', 'POST'])
def nonce_reuse():
    form = NonceReuseForm(meta={'csrf': False})
    key = nonce = None
    m1 = m2 = mf = c_forged = ''
    macs = None
    if form.is_submitted():
        key, nonce, m1, m2, mf = form.key.data, form.nonce.data, form.m1.data, form.m2.data, form.mf.data
        if form.validate():
            skey = binascii.unhexlify(key)
            snonce = binascii.unhexlify(nonce)
            c_forged, macs = solve(skey, snonce, bytes(m1, 'utf-8'), bytes(m2, 'utf-8'), bytes(mf, 'utf-8'))
    return render_template('nonce-reuse.html', form=form, key=key, nonce=nonce, m1=m1, m2=m2, mf=mf, c_forged=c_forged, macs=macs)

def solve(k, nonce, m1, m2, mf):
    aad1 = aad2 = b""
    c1, mac1 = gcm_encrypt(k, nonce, aad1, m1)
    c2, mac2 = gcm_encrypt(k, nonce, aad2, m2)

    possible_secrets = nonce_reuse_recover_secrets(nonce, aad1, aad2, c1, c2, mac1, mac2)
    c_forged = xor(c1, xor(m1, mf))
    aad_forged = b""
    macs = []
    for h, s in possible_secrets:
        mac = gmac(h, s, aad_forged, c_forged)
        macs.append((gf128_to_bytes(h), s, mac))
    return c_forged, macs