From bf9e2beda1c14305c05f6d776b07bc07a74b9bfd Mon Sep 17 00:00:00 2001
From: Alexander-D-Karpov <alexandr.d.karpov@gmail.com>
Date: Fri, 29 Dec 2023 00:24:01 +0300
Subject: [PATCH] added user option second factor totp

---
 akarpov/music/services/file.py           |   1 +
 akarpov/templates/base.html              |   5 +-
 akarpov/templates/users/disable_2fa.html |  24 +++++
 akarpov/templates/users/enable_2fa.html  |  22 +++++
 akarpov/templates/users/otp_verify.html  |  19 ++++
 akarpov/users/forms.py                   |   5 +
 akarpov/users/middleware.py              |  22 +++++
 akarpov/users/services/two_factor.py     |  14 +++
 akarpov/users/urls.py                    |   4 +
 akarpov/users/views.py                   | 112 ++++++++++++++++++++++-
 config/settings/base.py                  |  19 ++--
 config/urls.py                           |   2 +
 poetry.lock                              |  37 +++++++-
 pyproject.toml                           |   1 +
 14 files changed, 272 insertions(+), 15 deletions(-)
 create mode 100644 akarpov/templates/users/disable_2fa.html
 create mode 100644 akarpov/templates/users/enable_2fa.html
 create mode 100644 akarpov/templates/users/otp_verify.html
 create mode 100644 akarpov/users/services/two_factor.py

diff --git a/akarpov/music/services/file.py b/akarpov/music/services/file.py
index 1ae0359..3b4b9df 100644
--- a/akarpov/music/services/file.py
+++ b/akarpov/music/services/file.py
@@ -19,6 +19,7 @@ def load_dir(path: str, user_id: int):
 
 
 def load_file(path: str, user_id: int):
+    # TODO: convert to mp3 if not mp3
     process_mp3_file(path, user_id)
 
 
diff --git a/akarpov/templates/base.html b/akarpov/templates/base.html
index 00f300f..33a4ef7 100644
--- a/akarpov/templates/base.html
+++ b/akarpov/templates/base.html
@@ -93,8 +93,9 @@
                         <ul class="dropdown-menu dropdown-menu-dark text-small shadow" aria-labelledby="dropdownUser1">
                             <li><a class="dropdown-item {% active_link 'users:update' %}" href="{% url 'users:update'  %}">Settings</a></li>
                             <li><a class="dropdown-item {% active_link 'users:detail' request.user.username %}" href="{% url 'users:detail' request.user.username  %}">Profile</a></li>
-                            <li><a class="dropdown-item {% active_link 'tools:qr:create' %}" href="{% url 'tools:promocodes:activate' %}">Activate promocode</a></li>
-                            <li><a class="dropdown-item {% active_link 'tools:qr:create' %}" href="{% url 'users:history' %}">History</a></li>
+                            <li><a class="dropdown-item {% active_link 'tools:promocodes:activate' %}" href="{% url 'tools:promocodes:activate' %}">Activate promocode</a></li>
+                            <li><a class="dropdown-item {% active_link 'users:history' %}" href="{% url 'users:history' %}">History</a></li>
+                            <li><a class="dropdown-item {% active_link 'users:enable_2fa' %}" href="{% url 'users:enable_2fa' %}">2FA</a></li>
                             <li>
                                 <hr class="dropdown-divider">
                             </li>
diff --git a/akarpov/templates/users/disable_2fa.html b/akarpov/templates/users/disable_2fa.html
new file mode 100644
index 0000000..604678e
--- /dev/null
+++ b/akarpov/templates/users/disable_2fa.html
@@ -0,0 +1,24 @@
+{% extends "base.html" %}
+{% load crispy_forms_tags %}
+
+{% block content %}
+  <div class="container">
+    <div class="row">
+      <div class="col-md-6 mx-auto">
+        <div class="card mt-5">
+          <div class="card-body">
+            <h2 class="text-center">Disable Two-Factor Authentication</h2>
+            <p>If you want to disable two-factor authentication, please enter a token generated by your authentication app below. This is to verify your identity and to prevent unauthorized individuals from disabling your two-factor authentication.</p>
+            <p>If faced problems with disabling contact support</p>
+            <br>
+            <form method="POST">
+                {% csrf_token %}
+                {% crispy form %}
+                <input type="submit" class="btn btn-primary btn-block" value="Disable 2FA">
+            </form>
+          </div>
+        </div>
+      </div>
+    </div>
+  </div>
+{% endblock %}
diff --git a/akarpov/templates/users/enable_2fa.html b/akarpov/templates/users/enable_2fa.html
new file mode 100644
index 0000000..496ac3c
--- /dev/null
+++ b/akarpov/templates/users/enable_2fa.html
@@ -0,0 +1,22 @@
+{% extends "base.html" %}
+
+{% load crispy_forms_tags %}
+
+{% block content %}
+<div class="container">
+    <h1 class="mt-5">Enable Two-Factor Authentication</h1>
+    <p class="lead">Scan the QR code below or enter key with your TOTP app, then enter the code provided by the app to enable 2FA.</p>
+    <p>Your TOTP secret key is: {{ totp_key }}</p>
+
+    <div class="mb-4">
+        <img src="data:image/svg+xml;utf8,{{ qr_code_svg }}">
+    </div>
+
+    <form method="post" class="form">
+        {% csrf_token %}
+        {{ form|crispy }}
+        <button type="submit" class="btn btn-primary mt-2">Confirm</button>
+        <button type="submit" class="btn btn-secondary mt-2" name="cancel">Cancel</button>
+    </form>
+</div>
+{% endblock %}
diff --git a/akarpov/templates/users/otp_verify.html b/akarpov/templates/users/otp_verify.html
new file mode 100644
index 0000000..86e95d1
--- /dev/null
+++ b/akarpov/templates/users/otp_verify.html
@@ -0,0 +1,19 @@
+{% extends 'base.html' %}
+{% load crispy_forms_tags %}
+
+{% block content %}
+  <div class="row justify-content-center">
+    <div class="col-md-6">
+      <div class="card">
+        <div class="card-header bg-primary text-white">Enter OTP</div>
+        <div class="card-body">
+          <form method="post">
+            {% csrf_token %}
+            {{ form|crispy }}
+            <button type="submit" class="btn btn-primary mt-2">Submit OTP</button>
+          </form>
+        </div>
+      </div>
+    </div>
+  </div>
+{% endblock %}
diff --git a/akarpov/users/forms.py b/akarpov/users/forms.py
index 6e1dd9d..a61a853 100644
--- a/akarpov/users/forms.py
+++ b/akarpov/users/forms.py
@@ -1,5 +1,6 @@
 from allauth.account.forms import SignupForm
 from allauth.socialaccount.forms import SignupForm as SocialSignupForm
+from django import forms
 from django.contrib.auth import forms as admin_forms
 from django.contrib.auth import get_user_model
 from django.utils.translation import gettext_lazy as _
@@ -40,3 +41,7 @@ class UserSocialSignupForm(SocialSignupForm):
     Default fields will be added automatically.
     See UserSignupForm otherwise.
     """
+
+
+class OTPForm(forms.Form):
+    otp_token = forms.CharField()
diff --git a/akarpov/users/middleware.py b/akarpov/users/middleware.py
index 659c9c2..0b3258e 100644
--- a/akarpov/users/middleware.py
+++ b/akarpov/users/middleware.py
@@ -1,3 +1,5 @@
+from django.shortcuts import redirect
+from django.urls import resolve
 from django.utils.deprecation import MiddlewareMixin
 from rest_framework.exceptions import AuthenticationFailed
 
@@ -8,3 +10,23 @@ def process_request(self, request):
             if not request.user.is_verified:
                 raise AuthenticationFailed("Email is not verified")
         return None
+
+
+class Enforce2FAMiddleware:
+    def __init__(self, get_response):
+        self.get_response = get_response
+
+    def __call__(self, request):
+        response = self.get_response(request)
+
+        # Check user is authenticated and OTP token input is not completed
+        is_authenticated = request.user.is_authenticated
+        otp_not_verified = not request.session.get("otp_verified", False)
+        on_2fa_page = resolve(request.path_info).url_name == "enforce_otp_login"
+
+        # Enforce OTP token input, if user is authenticated but has not verified OTP, and is NOT on the 2FA page
+        if is_authenticated and otp_not_verified and not on_2fa_page:
+            request.session["next"] = request.get_full_path()
+            return redirect("users:enforce_otp_login")
+
+        return response
diff --git a/akarpov/users/services/two_factor.py b/akarpov/users/services/two_factor.py
new file mode 100644
index 0000000..7ce4ff2
--- /dev/null
+++ b/akarpov/users/services/two_factor.py
@@ -0,0 +1,14 @@
+import io
+
+import qrcode
+import qrcode.image.svg
+
+
+def generate_qr_code(data):
+    qr_image = qrcode.make(data, image_factory=qrcode.image.svg.SvgImage)
+    byte_arr = io.BytesIO()
+    qr_image.save(byte_arr)
+
+    qr_svg = byte_arr.getvalue().decode().replace("\n", "")
+
+    return qr_svg
diff --git a/akarpov/users/urls.py b/akarpov/users/urls.py
index 388d5ac..6f662c7 100644
--- a/akarpov/users/urls.py
+++ b/akarpov/users/urls.py
@@ -1,6 +1,8 @@
 from django.urls import include, path
 
 from akarpov.users.views import (
+    enable_2fa_view,
+    enforce_otp_login,
     user_detail_view,
     user_history_delete_view,
     user_history_view,
@@ -16,4 +18,6 @@
     path("history/", view=user_history_view, name="history"),
     path("history/delete", view=user_history_delete_view, name="history_delete"),
     path("<str:username>/", view=user_detail_view, name="detail"),
+    path("2fa/login", enforce_otp_login, name="enforce_otp_login"),
+    path("2fa/enable", enable_2fa_view, name="enable_2fa"),
 ]
diff --git a/akarpov/users/views.py b/akarpov/users/views.py
index 333fed3..c429754 100644
--- a/akarpov/users/views.py
+++ b/akarpov/users/views.py
@@ -1,12 +1,22 @@
+from allauth.account.views import LoginView
+from django.contrib import messages
 from django.contrib.auth import get_user_model
+from django.contrib.auth.decorators import login_required
 from django.contrib.auth.mixins import LoginRequiredMixin
 from django.contrib.messages.views import SuccessMessageMixin
-from django.urls import reverse
+from django.contrib.sites.shortcuts import get_current_site
+from django.http import HttpResponseRedirect
+from django.shortcuts import redirect, render
+from django.urls import reverse, reverse_lazy
 from django.utils.translation import gettext_lazy as _
 from django.views.generic import DetailView, ListView, RedirectView, UpdateView
+from django_otp import user_has_device
+from django_otp.plugins.otp_totp.models import TOTPDevice
 
+from akarpov.users.forms import OTPForm
 from akarpov.users.models import UserHistory
 from akarpov.users.services.history import create_history_warning_note
+from akarpov.users.services.two_factor import generate_qr_code
 from akarpov.users.themes.models import Theme
 
 User = get_user_model()
@@ -87,3 +97,103 @@ def get_redirect_url(self):
 
 
 user_history_delete_view = UserHistoryDeleteView.as_view()
+
+
+@login_required
+def enable_2fa_view(request):
+    user = request.user
+    devices = TOTPDevice.objects.filter(user=user, confirmed=True)
+
+    if devices.exists():
+        if request.method == "POST":
+            form = OTPForm(request.POST)
+            if form.is_valid():
+                token = form.cleaned_data["otp_token"]
+
+                # Verifying the token against all confirmed devices
+                for device in devices:
+                    if device.verify_token(token):
+                        device.delete()  # Delete the device if the token is valid
+
+                # Check if there are still confirmed devices left
+                if not TOTPDevice.objects.filter(user=user, confirmed=True).exists():
+                    messages.success(request, "Two-factor authentication disabled!")
+                    return redirect(reverse_lazy("blog:post_list"))
+                else:
+                    messages.error(request, "Invalid token, please try again.")
+            else:
+                messages.error(
+                    request, "The form submission was not valid. Please, try again."
+                )
+        form = OTPForm(initial={"otp_token": ""})
+        return render(request, "users/disable_2fa.html", {"form": form})
+
+    else:
+        if request.method == "POST":
+            device = TOTPDevice.objects.filter(user=user, confirmed=False).first()
+            if request.POST.get("cancel"):
+                device.delete()
+                return redirect(reverse_lazy("blog:post_list"))
+            form = OTPForm(request.POST)
+            if form.is_valid():
+                token = form.cleaned_data["otp_token"]
+                if device and device.verify_token(token):
+                    device.confirmed = True
+                    device.save()
+                    messages.success(request, "Two-factor authentication enabled!")
+                    return redirect(reverse_lazy("blog:post_list"))
+                else:
+                    messages.error(request, "Invalid token, please try again.")
+        else:
+            device, created = TOTPDevice.objects.get_or_create(
+                user=user, confirmed=False
+            )
+            site_name = get_current_site(request).name
+            provisioning_url = device.config_url
+            provisioning_url_site = provisioning_url.replace(
+                "otpauth://totp/", f"otpauth://totp/{site_name}:"
+            )
+            qr_code_svg = generate_qr_code(provisioning_url_site)
+            totp_key = device.key  # get device's secret key
+
+            form = OTPForm(initial={"otp_token": ""})
+
+        return render(
+            request,
+            "users/enable_2fa.html",
+            {"form": form, "qr_code_svg": qr_code_svg, "totp_key": totp_key},
+        )
+
+
+class OTPLoginView(LoginView):
+    def form_valid(self, form):
+        # Store URL for next page from the Login form
+        self.request.session["next"] = self.request.GET.get("next")
+
+        resp = super().form_valid(form)
+
+        # Redirect users with OTP devices to enter their tokens
+        if user_has_device(self.request.user):
+            return HttpResponseRedirect(reverse("users:enforce_otp_login"))
+
+        return resp
+
+
+login_view = OTPLoginView.as_view()
+
+
+@login_required
+def enforce_otp_login(request):
+    # TODO gather next url from loginrequired
+    if request.method == "POST":
+        form = OTPForm(request.POST)
+        if form.is_valid():
+            otp_token = form.cleaned_data["otp_token"]
+            device = TOTPDevice.objects.filter(user=request.user).first()
+            if device.verify_token(otp_token):
+                request.session["otp_verified"] = True
+                success_url = request.session.pop("next", None) or reverse_lazy("home")
+                return HttpResponseRedirect(success_url)
+    else:
+        form = OTPForm()
+    return render(request, "users/otp_verify.html", {"form": form})
diff --git a/config/settings/base.py b/config/settings/base.py
index 21e7a5a..7963a98 100644
--- a/config/settings/base.py
+++ b/config/settings/base.py
@@ -143,6 +143,12 @@
     "django_tables2",
     "location_field",
     "django_elasticsearch_dsl",
+    # 2fa
+    "django_otp",
+    "django_otp.plugins.otp_static",
+    "django_otp.plugins.otp_totp",
+    "django_otp.plugins.otp_hotp",
+    "django_otp.plugins.otp_email",
 ]
 
 HEALTH_CHECKS = [
@@ -237,6 +243,8 @@
     "django.middleware.common.CommonMiddleware",
     "django.middleware.csrf.CsrfViewMiddleware",
     "django.contrib.auth.middleware.AuthenticationMiddleware",
+    "django_otp.middleware.OTPMiddleware",
+    "akarpov.users.middleware.Enforce2FAMiddleware",
     "django.contrib.messages.middleware.MessageMiddleware",
     "django.middleware.common.BrokenLinkEmailsMiddleware",
     "django.middleware.clickjacking.XFrameOptionsMiddleware",
@@ -317,17 +325,6 @@
 
 # EMAIL
 # ------------------------------------------------------------------------------
-"""
-    host: EMAIL_HOST
-    port: EMAIL_PORT
-    username: EMAIL_HOST_USER
-    password: EMAIL_HOST_PASSWORD
-    use_tls: EMAIL_USE_TLS
-    use_ssl: EMAIL_USE_SSL
-    timeout: EMAIL_TIMEOUT
-    ssl_keyfile: EMAIL_SSL_KEYFILE
-    ssl_certfile: EMAIL_SSL_CERTFILE
-"""
 # https://docs.djangoproject.com/en/dev/ref/settings/#email-backend
 EMAIL_BACKEND = env(
     "DJANGO_EMAIL_BACKEND",
diff --git a/config/urls.py b/config/urls.py
index db10036..e63f2c6 100644
--- a/config/urls.py
+++ b/config/urls.py
@@ -14,6 +14,7 @@
 
 from akarpov.about.views import about_view, list_faq
 from akarpov.tools.shortener.views import redirect_view
+from akarpov.users.views import OTPLoginView
 from config.sitemaps import sitemaps
 
 urlpatterns = [
@@ -44,6 +45,7 @@
     path("gallery/", include("akarpov.gallery.urls", namespace="gallery")),
     path("ckeditor/", include("ckeditor_uploader.urls")),
     path("accounts/", include("allauth.urls")),
+    path("accounts/login/", OTPLoginView.as_view(), name="account_login"),
     path("", include("akarpov.blog.urls", namespace="blog")),
     path("<str:slug>", redirect_view, name="short_url"),
     # Your stuff: custom urls includes go here
diff --git a/poetry.lock b/poetry.lock
index 696b860..d65465f 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -4654,6 +4654,17 @@ files = [
     {file = "pyparsing-2.4.7.tar.gz", hash = "sha256:c203ec8783bf771a155b207279b9bccb8dea02d8f0c9e5f8ead507bc3246ecc1"},
 ]
 
+[[package]]
+name = "pypng"
+version = "0.20220715.0"
+description = "Pure Python library for saving and loading PNG images"
+optional = false
+python-versions = "*"
+files = [
+    {file = "pypng-0.20220715.0-py3-none-any.whl", hash = "sha256:4a43e969b8f5aaafb2a415536c1a8ec7e341cd6a3f957fd5b5f32a4cfeed902c"},
+    {file = "pypng-0.20220715.0.tar.gz", hash = "sha256:739c433ba96f078315de54c0db975aee537cbc3e1d0ae4ed9aab0ca1e427e2c1"},
+]
+
 [[package]]
 name = "pysocks"
 version = "1.7.1"
@@ -5005,6 +5016,30 @@ files = [
     {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"},
 ]
 
+[[package]]
+name = "qrcode"
+version = "7.4.2"
+description = "QR Code image generator"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "qrcode-7.4.2-py3-none-any.whl", hash = "sha256:581dca7a029bcb2deef5d01068e39093e80ef00b4a61098a2182eac59d01643a"},
+    {file = "qrcode-7.4.2.tar.gz", hash = "sha256:9dd969454827e127dbd93696b20747239e6d540e082937c90f14ac95b30f5845"},
+]
+
+[package.dependencies]
+colorama = {version = "*", markers = "platform_system == \"Windows\""}
+pillow = {version = ">=9.1.0", optional = true, markers = "extra == \"pil\""}
+pypng = "*"
+typing-extensions = "*"
+
+[package.extras]
+all = ["pillow (>=9.1.0)", "pytest", "pytest-cov", "tox", "zest.releaser[recommended]"]
+dev = ["pytest", "pytest-cov", "tox"]
+maintainer = ["zest.releaser[recommended]"]
+pil = ["pillow (>=9.1.0)"]
+test = ["coverage", "pytest"]
+
 [[package]]
 name = "rawpy"
 version = "0.19.0"
@@ -6838,4 +6873,4 @@ testing = ["coverage (>=5.0.3)", "zope.event", "zope.testing"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "5393f9258c796dbbf62840fdd2ba5d76e62f589f911ad767fbe093ba173b0ae8"
+content-hash = "d09c78af4c718b0fe6a1c0190627543ca8a902fc67611d688fb81603b93997fd"
diff --git a/pyproject.toml b/pyproject.toml
index f2221e0..b08ae38 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -115,6 +115,7 @@ numpy = "1.25.2"
 deep-translator = "1.4.2"
 textract = {git = "https://github.com/Alexander-D-Karpov/textract.git", branch = "master"}
 django-otp = "^1.3.0"
+qrcode = {extras = ["pil"], version = "^7.4.2"}
 
 
 [build-system]