key commitment

1 files changed, 78 insertions, 13 deletions

diff --git a/app.py b/app.py
index 9aade4e..800b263 100644
--- a/app.py
+++ b/app.py
@@ -1,12 +1,14 @@
-import binascii, secrets
+import binascii, secrets, io
 
-from flask import Flask, render_template, request, redirect, url_for
+from flask import Flask, render_template, request, redirect, url_for, send_file
 from flask_wtf import FlaskForm
-from wtforms import StringField
-from wtforms.validators import DataRequired, Length, ValidationError
+from flask_wtf.file import FileField, FileAllowed
+from werkzeug.datastructures import CombinedMultiDict
+from wtforms import StringField, RadioField
+from wtforms.validators import DataRequired, Length, ValidationError, InputRequired
 from Crypto.Cipher import AES
 
-from aesgcmanalysis import xor, gmac, gcm_encrypt, nonce_reuse_recover_secrets, gf128_to_bytes, nonce_truncation_recover_secrets
+from aesgcmanalysis import xor, gmac, gcm_encrypt, nonce_reuse_recover_secrets, gf128_to_bytes, mac_truncation_recover_secrets, att_merge_jpg_bmp
 
 app = Flask(__name__)
 
@@ -72,14 +74,14 @@ def solve_nonce_reuse(k, nonce, m1, m2, mf):
         macs.append((gf128_to_bytes(h), s, mac))
     return c_forged, macs
 
-class NonceTruncationForm(FlaskForm):
+class MACTruncationForm(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})
+@app.route('/mac-truncation', methods=['GET', 'POST'])
+def mac_truncation():
+    form = MACTruncationForm(meta={'csrf': False})
     key = nonce = None
     mf = ''
     h = c_forged = mac = None
@@ -88,11 +90,11 @@ def nonce_truncation():
         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,
+            h, c_forged, mac = solve_mac_truncation(skey, snonce, bytes(mf, 'utf-8'))
+    return render_template('mac-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):
+def solve_mac_truncation(k, nonce, mf):
     m = secrets.token_bytes(512)
     aad = b""
     c, mac = gcm_encrypt(k, nonce, aad, m, mac_bytes=1)
@@ -103,7 +105,70 @@ def solve_nonce_truncation(k, nonce, mf):
             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)
+        h, s = mac_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]
+
+class RequiredIf(InputRequired):
+    # adapted from https://stackoverflow.com/a/8464478
+    # a validator which makes a field required if
+    # another field is set and has a truthy value
+
+    def __init__(self, other_field_name, other_field_val, *args, **kwargs):
+        self.other_field_name = other_field_name
+        self.other_field_val = other_field_val
+        super(RequiredIf, self).__init__(*args, **kwargs)
+
+    def __call__(self, form, field):
+        other_field = form._fields.get(self.other_field_name)
+        if other_field is None:
+            raise Exception('no field named "%s" in form' % self.other_field_name)
+        if other_field.data == self.other_field_val:
+            super(RequiredIf, self).__call__(form, field)
+
+def FileSizeLimit(max_bytes, magic_start=None, magic_end=None):
+	# adapted from https://stackoverflow.com/a/67172432
+    def file_length_check(form, field):
+        if not field.data:
+            return
+        s = field.data.read()
+        n = len(field.data.read())
+        if n > max_bytes:
+            raise ValidationError(f"File size must be less than {max_bytes}B")
+        if magic_start and not s.startswith(magic_start):
+            raise ValidationError(f"Not a valid file; wrong initial magic bytes")
+        if magic_end and not s.endswith(magic_end):
+            raise ValidationError(f"Not a valid file; wrong final magic bytes")
+        field.data.seek(0)
+    return file_length_check
+
+class KeyCommitmentForm(FlaskForm):
+    mode = RadioField('mode', choices=[('sample', 'sample'), ('custom', 'custom')], validators=[DataRequired()])
+    jpeg = FileField('jpeg_file', validators=[FileAllowed(['jpg', 'jpeg']),
+											  RequiredIf('mode', 'custom'), FileSizeLimit(150000, b'\xff\xd8', b'\xff\xd9')])
+    bmp = FileField('bmp_file', validators=[FileAllowed(['bmp']),
+											RequiredIf('mode', 'custom'), FileSizeLimit(50000, b'\x42\x4d')])
+
+@app.route('/key-commitment', methods=['GET', 'POST'])
+def key_commitment():
+    form = KeyCommitmentForm(CombinedMultiDict((request.files, request.form)), meta={'csrf': False})
+    k1 = None
+    k2 = None
+    nonce = None
+    c = None
+
+    if form.is_submitted():
+        if form.validate():
+            if form.mode.data == 'sample':
+                jpeg_bytes = open('static/axolotl.jpg', 'rb').read()
+                bmp_bytes = open('static/kitten.bmp', 'rb').read()
+            else:
+                jpeg_bytes = form.jpeg.data.read()
+                bmp_bytes = form.bmp.data.read()
+            c, mac = att_merge_jpg_bmp(jpeg_bytes, bmp_bytes, aad=b"")
+            ct = c + mac
+            f = io.BytesIO(ct)
+            return send_file(f, mimetype='application/octet-stream', as_attachment=True, download_name="polyglot.enc")
+
+    return render_template('key-commitment.html', form=form, k1=k1, k2=k2, nonce=nonce, c=c)