trentpalmerdotorg/accounts/enable_totp.py

61 lines
2.3 KiB
Python

from django.shortcuts import redirect, render
import qrcode.image.svg
from .forms import EnableTotpForm
from django.contrib import messages
from .models import Account
from io import BytesIO
import pyotp
import qrcode
def disable_totp(request):
if not request.user.is_authenticated:
return redirect('audio:home')
if not request.user.account.use_totp:
return redirect('audio:home')
if request.method == "POST":
account = Account.objects.get(user=request.user)
account.use_totp = False
account.totp_key = None
account.save()
messages.success(request, 'Thanks for disabling 2fa!', extra_tags="mb-0")
return(redirect('accounts:edit_profile'))
return render(request, 'confirmation.html', {})
def enable_totp(request):
if not request.user.is_authenticated:
return redirect('audio:home')
if request.user.account.use_totp and request.user.account.totp_key is not None:
return(redirect('accounts:edit_profile'))
qr = get_totp_qr(request.user)
if request.method == "POST":
form = EnableTotpForm(request.POST, instance=request.user.account)
if form.is_valid():
totp_code = form.cleaned_data['totp_code']
if pyotp.TOTP(request.user.account.totp_key).verify(int(totp_code), valid_window=5):
account = Account.objects.get(user=request.user)
account.use_totp = True
account.save()
messages.success(request, 'Thanks for enabling 2fa!', extra_tags="mb-0")
return(redirect('accounts:edit_profile'))
else:
messages.error(request, 'Wrong Code, try again?', extra_tags="mb-0")
else:
form = EnableTotpForm(instance=request.user.account)
return render(request, 'accounts/totp_form.html', {'form': form, 'qr': qr})
def get_totp_qr(user):
if user.account.totp_key is None:
account = Account.objects.get(user=user)
account.totp_key = pyotp.random_base32()
account.save()
user.account.totp_key = account.totp_key
totp_uri = pyotp.totp.TOTP(user.account.totp_key).provisioning_uri(name='audio', issuer_name='trentpalmer.org')
img = qrcode.make(totp_uri, image_factory=qrcode.image.svg.SvgPathImage)
f = BytesIO()
img.save(f)
return(f.getvalue().decode('utf-8'))