mirror of
https://github.com/retspen/webvirtcloud
synced 2026-11-08 19:26:48 +00:00
72681ab54e
django-qr-code was used for one thing: the QR code in the OTP e-mail, as an inline SVG and a PNG. It brought pydantic, pydantic-core, annotated-types and typing-inspection (about 10 MB) and, from 5.0 on, Pillow. segno, which it used underneath, now draws both images with the same options (error correction M, 4-module border, 1.8 mm SVG modules, 12 px PNG modules); the SVG, the PNG and the alt text are byte for byte what django-qr-code 4.2 produced. account.html loaded the tag library without using it.
109 lines
5.1 KiB
Python
109 lines
5.1 KiB
Python
import base64
|
|
import hashlib
|
|
import re
|
|
import struct
|
|
import zlib
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from django.contrib.auth import get_user_model
|
|
from django.core import mail
|
|
from django.test import TestCase
|
|
from django_otp.plugins.otp_totp.models import TOTPDevice
|
|
from lxml import html
|
|
|
|
from accounts.utils import (get_existing_totp_device, get_user_totp_device, send_email_with_otp, validate_ssh_key)
|
|
|
|
|
|
class AccountsUtilsTestCase(TestCase):
|
|
def setUp(self):
|
|
User = get_user_model()
|
|
self.user = User.objects.create_user(
|
|
username="test_otp_user", email="user@example.com", password="password123"
|
|
)
|
|
|
|
def test_validate_ssh_key_valid(self):
|
|
blob = (
|
|
struct.pack(">I", 7)
|
|
+ b"ssh-rsa"
|
|
+ struct.pack(">I", 3)
|
|
+ b"\x01\x00\x01"
|
|
+ struct.pack(">I", 8)
|
|
+ b"\x00" * 8
|
|
)
|
|
b64 = base64.b64encode(blob).decode()
|
|
valid_key = f"ssh-rsa {b64} user@example.com"
|
|
self.assertTrue(validate_ssh_key(valid_key))
|
|
|
|
def test_validate_ssh_key_invalid_parts(self):
|
|
self.assertFalse(validate_ssh_key("only_two parts"))
|
|
self.assertFalse(validate_ssh_key("one_part"))
|
|
self.assertFalse(validate_ssh_key("four parts key here now"))
|
|
|
|
def test_validate_ssh_key_invalid_base64(self):
|
|
self.assertFalse(validate_ssh_key("ssh-rsa not-valid-b64! user@example.com"))
|
|
|
|
def test_validate_ssh_key_struct_error(self):
|
|
# Base64 with less than 4 bytes cannot unpack length
|
|
short_b64 = base64.b64encode(b"ab").decode()
|
|
self.assertFalse(validate_ssh_key(f"ssh-rsa {short_b64} user@example.com"))
|
|
|
|
def test_validate_ssh_key_type_mismatch(self):
|
|
blob = struct.pack(">I", 7) + b"ssh-rsa" + b"\x00" * 8
|
|
b64 = base64.b64encode(blob).decode()
|
|
# header says ssh-ed25519 but payload has ssh-rsa
|
|
self.assertFalse(validate_ssh_key(f"ssh-ed25519 {b64} user@example.com"))
|
|
|
|
def test_get_user_totp_device_creates_new(self):
|
|
self.assertEqual(TOTPDevice.objects.filter(user=self.user).count(), 0)
|
|
device = get_user_totp_device(self.user)
|
|
self.assertIsInstance(device, TOTPDevice)
|
|
self.assertEqual(TOTPDevice.objects.filter(user=self.user).count(), 1)
|
|
|
|
def test_get_user_totp_device_retrieves_existing(self):
|
|
device1 = get_user_totp_device(self.user)
|
|
device2 = get_user_totp_device(self.user)
|
|
self.assertEqual(device1.id, device2.id)
|
|
self.assertEqual(TOTPDevice.objects.filter(user=self.user).count(), 1)
|
|
|
|
def test_get_existing_totp_device_returns_none_without_a_device(self):
|
|
self.assertIsNone(get_existing_totp_device(self.user))
|
|
self.assertEqual(TOTPDevice.objects.filter(user=self.user).count(), 0)
|
|
|
|
def test_get_existing_totp_device_ignores_unconfirmed(self):
|
|
TOTPDevice.objects.create(user=self.user, confirmed=False)
|
|
self.assertIsNone(get_existing_totp_device(self.user))
|
|
|
|
def test_get_existing_totp_device_returns_confirmed(self):
|
|
device = TOTPDevice.objects.create(user=self.user, confirmed=True)
|
|
self.assertEqual(get_existing_totp_device(self.user).id, device.id)
|
|
|
|
def test_send_email_with_otp(self):
|
|
device = get_user_totp_device(self.user)
|
|
send_email_with_otp(self.user, device)
|
|
self.assertEqual(len(mail.outbox), 1)
|
|
self.assertIn("OTP QR Code", mail.outbox[0].subject)
|
|
self.assertEqual(mail.outbox[0].to, ["user@example.com"])
|
|
# the QR code as inline SVG and, for clients without SVG, as a PNG data URI
|
|
doc = html.fromstring(mail.outbox[0].alternatives[0][0])
|
|
self.assertEqual(len(doc.xpath("//svg")), 1)
|
|
(img,) = doc.xpath("//img")
|
|
self.assertEqual(img.get("alt"), device.config_url)
|
|
self.assertTrue(base64.b64decode(img.get("src").removeprefix("data:image/png;base64,")).startswith(b"\x89PNG"))
|
|
|
|
def test_otp_email_qr_codes_are_what_django_qr_code_drew(self):
|
|
# sha256 of the SVG and the PNG django-qr-code 4.2.0 made for this URL
|
|
url = "otpauth://totp/WebVirtCloud:alice?secret=JBSWY3DPEHPK3PXP&issuer=WebVirtCloud"
|
|
with patch("accounts.utils.send_mail") as send:
|
|
send_email_with_otp(self.user, SimpleNamespace(config_url=url))
|
|
page = send.call_args.kwargs["html_message"]
|
|
svg = re.search(r"<svg.*?</svg>", page, re.S).group(0)
|
|
png = base64.b64decode(re.search(r'<img src="data:image/png;base64,([^"]+)"', page).group(1))
|
|
self.assertEqual(hashlib.sha256(svg.encode()).hexdigest(), "b57323be725891af13d862a3a3351cac3ac64a493793c0258cf3e4a5886467c2")
|
|
# the PNG's header and pixel rows: its compressed bytes depend on the zlib build
|
|
chunks, pos = {}, 8
|
|
while pos < len(png):
|
|
size, kind = struct.unpack(">I4s", png[pos:pos + 8])
|
|
chunks[kind] = chunks.get(kind, b"") + png[pos + 8:pos + 8 + size]
|
|
pos += 12 + size
|
|
self.assertEqual(hashlib.sha256(chunks[b"IHDR"] + zlib.decompress(chunks[b"IDAT"])).hexdigest(), "cca16de611c73468827f70001f69dc2e48633e6aae2b570a9391eb40a1f8ae88")
|