Cloned DMOJ

judge/views/totp.py

@@ -0,0 +1,122 @@
+import base64
+from io import BytesIO
+
+import pyotp
+import qrcode
+from django.conf import settings
+from django.contrib.auth.mixins import LoginRequiredMixin
+from django.contrib.auth.views import SuccessURLAllowedHostsMixin
+from django.http import HttpResponseBadRequest, HttpResponseRedirect
+from django.urls import reverse
+from django.utils.http import is_safe_url
+from django.utils.translation import gettext as _
+from django.views.generic import FormView
+
+from judge.forms import TOTPForm
+from judge.utils.views import TitleMixin
+
+
+class TOTPView(TitleMixin, LoginRequiredMixin, FormView):
+    form_class = TOTPForm
+
+    def get_form_kwargs(self):
+        result = super(TOTPView, self).get_form_kwargs()
+        result['totp_key'] = self.profile.totp_key
+        return result
+
+    def dispatch(self, request, *args, **kwargs):
+        if request.user.is_authenticated:
+            self.profile = request.profile
+            if self.check_skip():
+                return self.next_page()
+        return super(TOTPView, self).dispatch(request, *args, **kwargs)
+
+    def check_skip(self):
+        raise NotImplementedError()
+
+    def next_page(self):
+        return HttpResponseRedirect(reverse('user_edit_profile'))
+
+
+class TOTPEnableView(TOTPView):
+    title = _('Enable Two Factor Authentication')
+    template_name = 'registration/totp_enable.html'
+
+    def get(self, request, *args, **kwargs):
+        profile = self.profile
+        if not profile.totp_key:
+            profile.totp_key = pyotp.random_base32(length=32)
+            profile.save()
+        return self.render_to_response(self.get_context_data())
+
+    def check_skip(self):
+        return self.profile.is_totp_enabled
+
+    def post(self, request, *args, **kwargs):
+        if not self.profile.totp_key:
+            return HttpResponseBadRequest('No TOTP key generated on server side?')
+        return super(TOTPEnableView, self).post(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['qr_code'] = self.render_qr_code(self.request.user.username, self.profile.totp_key)
+        return context
+
+    def form_valid(self, form):
+        self.profile.is_totp_enabled = True
+        self.profile.save()
+        # Make sure users don't get prompted to enter code right after enabling:
+        self.request.session['2fa_passed'] = True
+        return self.next_page()
+
+    @classmethod
+    def render_qr_code(cls, username, key):
+        totp = pyotp.TOTP(key)
+        uri = totp.provisioning_uri(username, settings.SITE_NAME)
+
+        qr = qrcode.QRCode(box_size=1)
+        qr.add_data(uri)
+        qr.make(fit=True)
+
+        image = qr.make_image(fill_color='black', back_color='white')
+        buf = BytesIO()
+        image.save(buf, format='PNG')
+        return 'data:image/png;base64,' + base64.b64encode(buf.getvalue()).decode('ascii')
+
+
+class TOTPDisableView(TOTPView):
+    title = _('Disable Two Factor Authentication')
+    template_name = 'registration/totp_disable.html'
+
+    def check_skip(self):
+        if not self.profile.is_totp_enabled:
+            return True
+        return settings.DMOJ_REQUIRE_STAFF_2FA and self.request.user.is_staff
+
+    def form_valid(self, form):
+        self.profile.is_totp_enabled = False
+        self.profile.totp_key = None
+        self.profile.save()
+        return self.next_page()
+
+
+class TOTPLoginView(SuccessURLAllowedHostsMixin, TOTPView):
+    title = _('Perform Two Factor Authentication')
+    template_name = 'registration/totp_auth.html'
+
+    def check_skip(self):
+        return not self.profile.is_totp_enabled or self.request.session.get('2fa_passed', False)
+
+    def next_page(self):
+        redirect_to = self.request.GET.get('next', '')
+        url_is_safe = is_safe_url(
+            url=redirect_to,
+            allowed_hosts=self.get_success_url_allowed_hosts(),
+            require_https=self.request.is_secure(),
+        )
+        return HttpResponseRedirect((redirect_to if url_is_safe else '') or reverse('user_page'))
+
+    def form_valid(self, form):
+        self.request.session['2fa_passed'] = True
+        return self.next_page()