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

dmoj/urls.py

@@ -85,6 +85,8 @@
     url(r'^2fa/enable/$', totp.TOTPEnableView.as_view(), name='enable_2fa'),
     url(r'^2fa/disable/$', totp.TOTPDisableView.as_view(), name='disable_2fa'),
 
+    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/forms.py

@@ -1,5 +1,7 @@
+import json
 from operator import attrgetter
 
+
 import pyotp
 from django import forms
 from django.conf import settings
@@ -129,14 +131,30 @@ class TOTPForm(Form):
     totp_token = NoAutoCompleteCharField(validators=[
         RegexValidator('^[0-9]{6}$', _('Two Factor Authentication tokens must be 6 decimal digits.')),
     ])
+    scratch_code = NoAutoCompleteCharField(validators=[
+        RegexValidator('^[A-Z2-7]{16}$', _('Scratch codes must be 16 decimal digits.')),
+    ])
 
     def __init__(self, *args, **kwargs):
         self.totp_key = kwargs.pop('totp_key')
+        self.scratch_codes = json.loads(kwargs.pop('scratch_codes'))
+        self.profile = kwargs.pop('profile')
         super(TOTPForm, self).__init__(*args, **kwargs)
+        self.fields['totp_token'].required = False
+        self.fields['scratch_code'].required = False
 
     def clean_totp_token(self):
-        if not pyotp.TOTP(self.totp_key).verify(self.cleaned_data['totp_token'], valid_window=self.TOLERANCE):
+        totp_token = self.cleaned_data.get('totp_token')
+        scratch_code = self.data.get('scratch_code')
+        if totp_token and scratch_code:
+            raise ValidationError(_('Enter only the 6-digit code or one of your scratch codes'))
+        if not scratch_code and not pyotp.TOTP(self.totp_key).verify(totp_token, valid_window=self.TOLERANCE):
             raise ValidationError(_('Invalid Two Factor Authentication token.'))
+        if not totp_token and scratch_code in self.scratch_codes:
+            self.profile.scratch_codes = json.dumps(self.scratch_codes.remove(scratch_code))
+            self.profile.save()
+        elif scratch_code:
+            raise ValidationError(_('Invalid Scratch Code.'))
 
 
 class ProblemCloneForm(Form):

judge/migrations/0103_scratch_codes.py

@@ -0,0 +1,22 @@
+# Generated by Django 2.2.7 on 2020-04-23 03:28
+
+import django.core.validators
+import django.db.models.deletion
+from django.db import migrations
+
+import judge.models.profile
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('judge', '0102_api_token'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='profile',
+            name='scratch_codes',
+            field=judge.models.profile.EncryptedNullCharField(blank=True, help_text='JSON of 16 character base32-encoded codes for scratch codes', max_length=255, null=True, validators=[django.core.validators.RegexValidator('^[A-Z2-7]$', 'Scratch codes must be empty or base32')], verbose_name='Scratch Codes'),
+        ),
+    ]

judge/models/profile.py

@@ -1,7 +1,9 @@
 import base64
 import hmac
+import json
 import secrets
 import struct
+import pyotp
 from operator import mul
 
 from django.conf import settings
@@ -123,6 +125,10 @@ class Profile(models.Model):
                                                             _('API token must be None or hexadecimal'))])
     notes = models.TextField(verbose_name=_('internal notes'), null=True, blank=True,
                              help_text=_('Notes for administrators regarding this user.'))
+    scratch_codes = EncryptedNullCharField(max_length=255, null=True, blank=True, verbose_name=_('Scratch Codes'),
+                                           help_text=_('JSON of 16 character base32-encoded codes for scratch codes'),
+                                           validators=[RegexValidator('^[A-Z2-7]$',
+                                                                      _('Scratch codes must be empty or base32'))])
 
     @cached_property
     def organization(self):
@@ -170,6 +176,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(5)]
+        self.scratch_codes = json.dumps(codes)
+        self.save(update_fields=['scratch_codes'])
+        return '\n'.join(codes)
+
+    generate_scratch_codes.alters_data = True
+
     def remove_contest(self):
         self.current_contest = None
         self.save()

judge/views/totp.py

@@ -1,6 +1,8 @@
 import base64
+import json
 from io import BytesIO
 
+
 import pyotp
 import qrcode
 from django.conf import settings
@@ -22,6 +24,8 @@ class TOTPView(TitleMixin, LoginRequiredMixin, FormView):
     def get_form_kwargs(self):
         result = super(TOTPView, self).get_form_kwargs()
         result['totp_key'] = self.profile.totp_key
+        result['scratch_codes'] = self.profile.scratch_codes
+        result['profile'] = self.profile
         return result
 
     def dispatch(self, request, *args, **kwargs):
@@ -47,6 +51,9 @@ def get(self, request, *args, **kwargs):
         if not profile.totp_key:
             profile.totp_key = pyotp.random_base32(length=32)
             profile.save()
+        if not profile.scratch_codes:
+            profile.scratch_codes = json.dumps([pyotp.random_base32(length=16) for i in range(5)])
+            profile.save()
         return self.render_to_response(self.get_context_data())
 
     def check_skip(self):
@@ -60,6 +67,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'] = '\n'.join(json.loads(self.profile.scratch_codes))
         context['qr_code'] = self.render_qr_code(self.request.user.username, self.profile.totp_key)
         return context
 
@@ -97,6 +105,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

@@ -281,6 +281,28 @@ 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()}})
+
+
+@require_POST
+@login_required
+def remove_scratch_codes(request):
+    profile = request.profile
+    with transaction.atomic(), revisions.create_revision():
+        profile.scratch_codes = None
+        profile.save()
+        revisions.set_user(request.user)
+        revisions.set_comment(_('Removed scratch codes for user'))
+    return JsonResponse({})
+
+
 class UserList(QueryStringSortMixin, DiggPaginatorMixin, TitleMixin, ListView):
     model = Profile
     title = gettext_lazy('Leaderboard')

templates/registration/totp_auth.html

@@ -38,6 +38,6 @@
             <hr>
             <button style="float:right;" type="submit">{{ _('Login!') }}</button>
         </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-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 %}

templates/registration/totp_disable.html

@@ -45,6 +45,10 @@
                 <label class="inline-header grayed">{{ _('Enter the 6-digit code generated by your app:') }}</label>
                 <span class="fullwidth">{{ form.totp_token }}</span>
             </div>
+            <div>
+                <label class="inline-header grayed">{{ _('Or enter one of your scratch codes:') }}</label>
+                <span class="fullwidth">{{ form.scratch_code }}</span>
+            </div>
             <button style="margin-top: 0.5em" class="button" type="submit">{{ _('Disable Two Factor Authentication') }}</button>
         </form>
     </div>

templates/registration/totp_enable.html

@@ -42,6 +42,10 @@
             image-rendering: -moz-crisp-edges;
             image-rendering: crisp-edges;
         }
+
+        .scratch-codes {
+            white-space: pre-line;
+        }
     </style>
 {% endblock %}
 
@@ -61,6 +65,52 @@
             }
         });
     </script>
+    <script src="{{ static('libs/clipboard/clipboard.js') }}"></script>
+    <script src="{{ static('libs/clipboard/tooltip.js') }}"></script>
+    <script type="text/javascript">
+        $(function () {
+            var info_float = $('.info-float');
+            if (info_float.length) {
+                var container = $('#content-right');
+                if (window.bad_browser) {
+                    container.css('float', 'right');
+                } else if (!featureTest('position', 'sticky')) {
+                    fix_div(info_float, 55);
+                    $(window).resize(function () {
+                        info_float.width(container.width());
+                    });
+                    info_float.width(container.width());
+                }
+            }
+
+            var copyButton;
+            $('pre code').each(function () {
+                $(this).parent().before($('<div>', {'class': 'copy-clipboard'})
+                    .append(copyButton = $('<span>', {
+                        'class': 'btn-clipboard',
+                        'data-clipboard-text': $(this).text(),
+                        'title': 'Click to copy'
+                    }).text('Copy')));
+
+                $(copyButton.get(0)).mouseleave(function () {
+                    $(this).attr('class', 'btn-clipboard');
+                    $(this).removeAttr('aria-label');
+                });
+
+                var curClipboard = new Clipboard(copyButton.get(0));
+
+                curClipboard.on('success', function (e) {
+                    e.clearSelection();
+                    showTooltip(e.trigger, 'Copied!');
+                });
+
+                curClipboard.on('error', function (e) {
+                    showTooltip(e.trigger, fallbackMessage(e.action));
+                });
+
+            });
+        });
+    </script>
 {% endblock %}
 
 {% block body %}
@@ -84,6 +134,15 @@
                 <span class="fullwidth">{{ form.totp_token }}</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 }}</span></code></pre>
+            </div>
         </form>
     </div>
 {% endblock %}