import binascii, secrets

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 Crypto.Cipher import AES

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

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):
        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_nonce_reuse(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_nonce_reuse(k, nonce, m1, m2, mf):
    aad1 = aad2 = b""
    c1, mac1 = gcm_encrypt(k, nonce, aad1, m1)
    c2, mac2 = gcm_encrypt(k, nonce, aad2, m2)

    default_m1 = 'The universe (which others call the Library)'
    default_m2 = 'From any of the hexagons one can see, interminably'
    if k == b'tlonorbistertius' and nonce == b'JORGELBORGES' and m1 == default_m1 and m2 == default_m2:
        possible_secrets = [(144676297626548424623350164317265032260,
                             137128696435097309357166918744288944691),
                            (176085395972970454284981815262084281580,
                             250035608282660492164551282952970544944)]
    else:
        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

class NonceTruncationForm(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])
    mf = StringField('forged message', validators=[DataRequired(), Length(min=1, max=64)])

@app.route('/nonce-truncation', methods=['GET', 'POST'])
def nonce_truncation():
    form = NonceTruncationForm(meta={'csrf': False})
    key = nonce = None
    mf = ''
    h = c_forged = mac = None
    if form.is_submitted():
        key, nonce, mf = form.key.data, form.nonce.data, form.mf.data
        if form.validate():
            skey = binascii.unhexlify(key)
            snonce = binascii.unhexlify(nonce)
            h, c_forged, mac = solve_nonce_truncation(skey, snonce, bytes(mf, 'utf-8'))
    return render_template('nonce-truncation.html', form=form, key=key, nonce=nonce,
                           mf=mf, h=h, c_forged=c_forged, mac=mac)

def solve_nonce_truncation(k, nonce, mf):
    m = secrets.token_bytes(512)
    aad = b""
    c, mac = gcm_encrypt(k, nonce, aad, m, mac_bytes=1)
    if k == b'tlonorbistertius' and nonce == b'JORGELBORGES':
        h, s = 176085395972970454284981815262084281580, 48
    else:
        def oracle(base, aad, mac, nonce):
            cipher = AES.new(k, mode=AES.MODE_GCM, nonce=nonce, mac_len=1)
            cipher.update(aad)
            cipher.decrypt_and_verify(base, mac)
        h, s = nonce_truncation_recover_secrets(c, mac, nonce, 1, aad, oracle)
    c_forged, aad_forged = xor(c, xor(m, mf)), b""
    mac = gmac(h, s, aad_forged, c_forged)
    return gf128_to_bytes(h), c_forged, mac[:1]