trentpalmerdotorg/accounts/enable_totp.py

61 lines
2.3 KiB
Python
Raw Normal View History

2021-02-24 20:13:54 -08:00
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'))
2021-02-24 20:13:54 -08:00
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'))