Add scratch code support for 2FA; fixes DMOJ#784

dmoj/settings.py

@@ -73,6 +73,7 @@
 DMOJ_BLOG_NEW_PROBLEM_COUNT = 7
 DMOJ_BLOG_RECENTLY_ATTEMPTED_PROBLEMS_COUNT = 7
 DMOJ_TOTP_TOLERANCE_HALF_MINUTES = 1
+DMOJ_SCRATCH_CODES_COUNT = 5
 DMOJ_USER_MAX_ORGANIZATION_COUNT = 3
 # Whether to allow users to download their data
 DMOJ_USER_DATA_DOWNLOAD = False

dmoj/urls.py

@@ -80,6 +80,8 @@
     url(r'^2fa/webauthn/assert/$', two_factor.WebAuthnAttestView.as_view(), name='webauthn_assert'),
     url(r'^2fa/webauthn/delete/(?P<pk>\d+)$', two_factor.WebAuthnDeleteView.as_view(), name='webauthn_delete'),
 
+    url(r'^2fa/scratchcode/generate/$', user.generate_scratch_codes, name='generate_scratch_codes'),
+
     url(r'api/token/generate/$', user.generate_api_token, name='generate_api_token'),
     url(r'api/token/remove/$', user.remove_api_token, name='remove_api_token'),
 ]

judge/admin/profile.py

@@ -66,6 +66,7 @@ def get_fields(self, request, obj=None):
         if request.user.has_perm('judge.totp'):
             fields = list(self.fields)
             fields.insert(fields.index('is_totp_enabled') + 1, 'totp_key')
+            fields.insert(fields.index('totp_key') + 1, 'scratch_codes')
             return tuple(fields)
         else:
             return self.fields

judge/forms.py

@@ -20,6 +20,21 @@
 from judge.widgets import HeavyPreviewPageDownWidget, MathJaxPagedownWidget, PagedownWidget, Select2MultipleWidget, \
     Select2Widget
 
+totp_or_scratch_code_validators = {
+    6: {
+        'regex_validator': RegexValidator('^[0-9]{6}$',
+                                          _('Two-factor authentication tokens must be 6 decimal digits.')),
+        'verify': lambda code, profile: not pyotp.TOTP(profile.totp_key)
+                                                 .verify(code, valid_window=settings.DMOJ_TOTP_TOLERANCE_HALF_MINUTES),
+        'err': ValidationError(_('Invalid two-factor authentication token.')),
+    },
+    16: {
+        'regex_validator': RegexValidator('^[A-Z0-9]{16}$', _('Scratch codes must be 16 base32 characters.')),
+        'verify': lambda code, profile: code not in json.loads(profile.scratch_codes),
+        'err': ValidationError(_('Invalid scratch code.')),
+    },
+}
+
 
 def fix_unicode(string, unsafe=tuple('\u202a\u202b\u202d\u202e')):
     return string + (sum(k in unsafe for k in string) - string.count('\u202c')) * '\u202c'
@@ -164,20 +179,23 @@ def widget_attrs(self, widget):
 class TOTPForm(Form):
     TOLERANCE = settings.DMOJ_TOTP_TOLERANCE_HALF_MINUTES
 
-    totp_token = NoAutoCompleteCharField(validators=[
-        RegexValidator('^[0-9]{6}$', _('Two-factor authentication tokens must be 6 decimal digits.')),
-    ], required=False)
+    totp_or_scratch_code = NoAutoCompleteCharField(required=False)
     webauthn_response = forms.CharField(widget=forms.HiddenInput(), required=False)
 
     def __init__(self, *args, **kwargs):
         self.profile = kwargs.pop('profile')
         super().__init__(*args, **kwargs)
 
     def clean(self):
-        if (not self.cleaned_data.get('totp_token') or
-                not pyotp.TOTP(self.profile.totp_key).verify(self.cleaned_data['totp_token'],
-                                                             valid_window=self.TOLERANCE)):
-            raise ValidationError(_('Invalid two-factor authentication token.'))
+        totp_or_scratch_code = self.cleaned_data.get('totp_or_scratch_code')
+        totp_or_scratch_code_len = len(totp_or_scratch_code)
+        if totp_or_scratch_code_len not in totp_or_scratch_code_validators:
+            raise ValidationError(_('Invalid code length.'))
+        else:
+            validator = totp_or_scratch_code_validators[totp_or_scratch_code_len]
+            validator['regex_validator'](totp_or_scratch_code)
+            if validator['verify'](totp_or_scratch_code, self.profile):
+                raise validator['err']
 
 
 class TwoFactorLoginForm(TOTPForm):
@@ -189,6 +207,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     def clean(self):
+        totp_or_scratch_code = self.cleaned_data.get('totp_or_scratch_code')
         if self.profile.is_webauthn_enabled and self.cleaned_data.get('webauthn_response'):
             if len(self.cleaned_data['webauthn_response']) > 65536:
                 raise ValidationError(_('Invalid WebAuthn response.'))
@@ -220,10 +239,16 @@ def clean(self):
 
             credential.counter = sign_count
             credential.save(update_fields=['counter'])
-        elif self.profile.is_totp_enabled and self.cleaned_data.get('totp_token'):
-            if pyotp.TOTP(self.profile.totp_key).verify(self.cleaned_data['totp_token'], valid_window=self.TOLERANCE):
+        elif self.profile.is_totp_enabled and totp_or_scratch_code:
+            if pyotp.TOTP(self.profile.totp_key).verify(totp_or_scratch_code, valid_window=self.TOLERANCE):
+                return
+            elif totp_or_scratch_code in json.loads(self.profile.scratch_codes):
+                scratch_codes = json.loads(self.profile.scratch_codes)
+                scratch_codes.remove(totp_or_scratch_code)
+                self.profile.scratch_codes = json.dumps(scratch_codes)
+                self.profile.save()
                 return
-            raise ValidationError(_('Invalid two-factor authentication token.'))
+            raise ValidationError(_('Invalid two-factor authentication token or scratch code.'))
         else:
             raise ValidationError(_('Must specify either totp_token or webauthn_response.'))
 

judge/migrations/0109_scratch_codes.py

@@ -0,0 +1,22 @@
+# Generated by Django 2.2.7 on 2020-05-27 05:50
+
+import django.core.validators
+import django.db.models.deletion
+from django.db import migrations
+
+import judge.models.profile
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('judge', '0108_bleach_problems'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='profile',
+            name='scratch_codes',
+            field=judge.models.profile.EncryptedNullCharField(blank=True, help_text='JSON array of 16 character base32-encoded codes for scratch codes', max_length=255, null=True, validators=[django.core.validators.RegexValidator(r'^["\[\],A-Z0-9 ]*$', 'Scratch codes must be empty or a JSON array of 16-character base32 codes')], verbose_name='scratch codes'),
+        ),
+    ]

judge/models/profile.py

@@ -1,9 +1,11 @@
 import base64
 import hmac
+import json
 import secrets
 import struct
 from operator import mul
 
+import pyotp
 import webauthn
 from django.conf import settings
 from django.contrib.auth.models import User
@@ -121,6 +123,11 @@ class Profile(models.Model):
                                       help_text=_('32 character base32-encoded key for TOTP'),
                                       validators=[RegexValidator('^$|^[A-Z2-7]{32}$',
                                                                  _('TOTP key must be empty or base32'))])
+    scratch_codes = EncryptedNullCharField(max_length=255, null=True, blank=True, verbose_name=_('scratch codes'),
+                                           help_text=_('JSON array of 16 character base32-encoded codes for scratch codes'),
+                                           validators=[RegexValidator(r'^["\[\],A-Z0-9 ]*$',
+                                                                      _('Scratch codes must be empty or a JSON array of \
+                                                                         16-character base32 codes'))])
     api_token = models.CharField(max_length=64, null=True, verbose_name=_('API token'),
                                  help_text=_('64 character hex-encoded API access token'),
                                  validators=[RegexValidator('^[a-f0-9]{64}$',
@@ -175,6 +182,14 @@ def generate_api_token(self):
 
     generate_api_token.alters_data = True
 
+    def generate_scratch_codes(self):
+        codes = [pyotp.random_base32(length=16) for i in range(settings.DMOJ_SCRATCH_CODES_COUNT)]
+        self.scratch_codes = json.dumps(codes)
+        self.save(update_fields=['scratch_codes'])
+        return codes
+
+    generate_scratch_codes.alters_data = True
+
     def remove_contest(self):
         self.current_contest = None
         self.save()

judge/views/two_factor.py

@@ -53,7 +53,9 @@ def get(self, request, *args, **kwargs):
         profile = self.profile
         if not profile.totp_key:
             profile.totp_key = pyotp.random_base32(length=32)
-            profile.save()
+            profile.save(update_fields=['totp_key'])
+        if not profile.scratch_codes:
+            profile.generate_scratch_codes()
         return self.render_to_response(self.get_context_data())
 
     def check_skip(self):
@@ -67,6 +69,7 @@ def post(self, request, *args, **kwargs):
     def get_context_data(self, **kwargs):
         context = super(TOTPEnableView, self).get_context_data(**kwargs)
         context['totp_key'] = self.profile.totp_key
+        context['scratch_codes'] = json.loads(self.profile.scratch_codes)
         context['qr_code'] = self.render_qr_code(self.request.user.username, self.profile.totp_key)
         return context
 
@@ -104,6 +107,7 @@ def check_skip(self):
     def form_valid(self, form):
         self.profile.is_totp_enabled = False
         self.profile.totp_key = None
+        self.profile.scratch_codes = None
         self.profile.save()
         return self.next_page()
 

judge/views/user.py

@@ -408,6 +408,16 @@ def remove_api_token(request):
     return JsonResponse({})
 
 
+@require_POST
+@login_required
+def generate_scratch_codes(request):
+    profile = request.profile
+    with transaction.atomic(), revisions.create_revision():
+        revisions.set_user(request.user)
+        revisions.set_comment(_('Generated scratch codes for user'))
+    return JsonResponse({'data': {'codes': profile.generate_scratch_codes()}})
+
+
 class UserList(QueryStringSortMixin, DiggPaginatorMixin, TitleMixin, ListView):
     model = Profile
     title = gettext_lazy('Leaderboard')

templates/common-content.html

@@ -40,12 +40,12 @@
 
                     var curClipboard = new Clipboard(copyButton.get(0));
 
-                    curClipboard.on('success', function (e) {
+                    curClipboard.on('success', function(e) {
                         e.clearSelection();
                         showTooltip(e.trigger, 'Copied!');
                     });
 
-                    curClipboard.on('error', function (e) {
+                    curClipboard.on('error', function(e) {
                         showTooltip(e.trigger, fallbackMessage(e.action));
                     });
 

templates/registration/totp_disable.html

@@ -36,14 +36,14 @@
         <form id="totp-disable-form" action="" method="post" class="form-area">
             {% csrf_token %}
             <div class="block-header">{{ _('To protect your account, you must first authenticate before you can disable Two Factor Authentication.') }}</div>
-            {% if form.totp_token.errors %}
+            {% if form.errors %}
                 <div class="form-errors">
-                    {{ form.totp_token.errors }}
+                    {{ form.non_field_errors() }}
                 </div>
             {% endif %}
             <div>
-                <label class="inline-header grayed">{{ _('Enter the 6-digit code generated by your app:') }}</label>
-                <span class="fullwidth">{{ form.totp_token }}</span>
+                <label class="inline-header grayed">{{ _('Enter the 6-digit code generated by your app or one of your 16-digit scratch codes:') }}</label>
+                <span class="fullwidth">{{ form.totp_or_scratch_code }}</span>
             </div>
             <button style="margin-top: 0.5em" class="button" type="submit">{{ _('Disable Two Factor Authentication') }}</button>
         </form>

templates/registration/totp_enable.html

@@ -1,4 +1,4 @@
-{% extends "base.html" %}
+{% extends "common-content.html" %}
 
 {% block media %}
     <style>
@@ -42,10 +42,14 @@
             image-rendering: -moz-crisp-edges;
             image-rendering: crisp-edges;
         }
+
+        .scratch-codes {
+            white-space: pre-line;
+        }
     </style>
 {% endblock %}
 
-{% block js_media %}
+{% block content_js_media %}
     <script>
         $(function () {
             if (window.navigator.userAgent.indexOf('Edge') >= 0) {
@@ -74,17 +78,26 @@
                 <code class="totp-code">{{ totp_key }}</code>
             </div>
             <hr>
-            {% if form.totp_token.errors or form.non_field_errors() %}
+            {% if form.totp_or_scratch_code.errors or form.non_field_errors() %}
                 <div class="form-errors">
-                    {{ form.totp_token.errors }}
+                    {{ form.totp_or_scratch_code.errors }}
                     {{ form.non_field_errors() }}
                 </div>
             {% endif %}
             <div>
                 <label class="inline-header grayed">{{ _('Enter the 6-digit code generated by your app:') }}</label>
-                <span class="fullwidth">{{ form.totp_token }}</span>
+                <span class="fullwidth">{{ form.totp_or_scratch_code }}</span>
             </div>
             <button style="margin-top: 0.5em" class="button" type="submit">{{ _('Enable Two Factor Authentication') }}</button>
+            <div class="scratch-codes">
+                <label class="inline-header grayed">{{ _('Below is a list of one-time use scratch codes. 
+                                                          These codes can only be used once and are for emergency use.
+                                                          You can use these codes to login to your account or disable two-factor authentication.
+                                                          If you ever need more scratch codes, you can regenerate them on the edit profile tab.
+                                                          Please write these down and keep them in a secure location.
+                                                          These codes will never be shown again after you enable two-factor authentication.') }}</label>
+                <pre><code><span class="fullwidth">{{ scratch_codes|join('\n') }}</span></code></pre>
+            </div>
         </form>
     </div>
 {% endblock %}

templates/registration/two_factor_auth.html

@@ -9,15 +9,15 @@
             left: 50%;
         }
 
-        #totp-token-container {
+        #totp-or-scratch-code-container {
             margin: 0.5em 0;
         }
 
         #id_totp_token {
             width: 100%;
         }
 
-        .totp-panel-message {
+        .totp-or-scratch-code-panel-message {
             width: 300px;
         }
     </style>
@@ -81,8 +81,8 @@
             {% endif %}
 
             {% if request.profile.is_totp_enabled %}
-                <div><label class="inline-header grayed">{{ _('Enter the 6-digit code generated by your app:') }}</label></div>
-                <div id="totp-token-container"><span class="fullwidth">{{ form.totp_token }}</span></div>
+                <div><label class="inline-header grayed">{{ _('Enter the 6-digit code generated by your app <br> or one of your 16-character scratch codes:') }}</label></div>
+                <div id="totp-or-scratch-code-container"><span class="fullwidth">{{ form.totp_or_scratch_code }}</span></div>
                 <hr>
             {% endif %}
             {% if request.profile.is_webauthn_enabled %}
@@ -93,6 +93,6 @@
                 <button style="float:right;" type="submit">{{ _('Login!') }}</button>
             {% endif %}
         </form>
-        <p class="totp-panel-message">{{ _('If you lost your authentication device, please contact us at %(email)s.', email=SITE_ADMIN_EMAIL)|urlize }}</p>
+        <p class="totp-or-scratch-code-panel-message">{{ _('If you lost your authentication device and are unable to use your scratch codes, please contact us at %(email)s.', email=SITE_ADMIN_EMAIL)|urlize }}</p>
     </div>
 {% endblock %}