From feb17e8071e2d8449fd1498d2d0caab300a764df Mon Sep 17 00:00:00 2001
From: David Guillot <david@guillot.me>
Date: Thu, 12 Dec 2024 17:00:33 +0100
Subject: [PATCH] feat: Django command to generate/export a Vapid keypair

---
 README.md                                     |  6 ++++-
 webpush/management/__init__.py                |  0
 webpush/management/commands/__init__.py       |  0
 .../webpush_generate_vapid_keypair.py         | 10 ++++++++
 webpush/tests.py                              |  3 ---
 webpush/tests/__init__.py                     |  0
 webpush/tests/test_vapid.py                   | 25 +++++++++++++++++++
 webpush/vapid.py                              | 16 ++++++++++++
 8 files changed, 56 insertions(+), 4 deletions(-)
 create mode 100644 webpush/management/__init__.py
 create mode 100644 webpush/management/commands/__init__.py
 create mode 100644 webpush/management/commands/webpush_generate_vapid_keypair.py
 delete mode 100644 webpush/tests.py
 create mode 100644 webpush/tests/__init__.py
 create mode 100644 webpush/tests/test_vapid.py
 create mode 100644 webpush/vapid.py

diff --git a/README.md b/README.md
index 97faad1..e3bfee4 100644
--- a/README.md
+++ b/README.md
@@ -36,7 +36,11 @@ WEBPUSH_SETTINGS = {
 ```
 **Replace ``"Vapid Public Key"`` and ``"Vapid Private Key"`` with your Vapid Keys. Also replace ``admin@example.com`` with your email so that the push server of browser can reach to you if anything goes wrong.**
 
-> **To know how to obtain Vapid Keys please see this [`py_vapid`](https://github.com/web-push-libs/vapid/tree/master/python) and [Google Developer Documentation](https://developers.google.com/web/fundamentals/push-notifications/subscribing-a-user#how_to_create_application_server_keys). You can obtain one easily from [web-push-codelab.glitch.me](https://web-push-codelab.glitch.me/). ``Application Server Keys`` and ``Vapid Keys`` both are same.**
+**Generate a Vapid key pair**
+
+```shell
+python manage.py webpush_generate_vapid_keypair
+```
 
 Then include `webpush` in the `urls.py`
 
diff --git a/webpush/management/__init__.py b/webpush/management/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/webpush/management/commands/__init__.py b/webpush/management/commands/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/webpush/management/commands/webpush_generate_vapid_keypair.py b/webpush/management/commands/webpush_generate_vapid_keypair.py
new file mode 100644
index 0000000..accf799
--- /dev/null
+++ b/webpush/management/commands/webpush_generate_vapid_keypair.py
@@ -0,0 +1,10 @@
+from django.core.management.base import BaseCommand
+
+from ...vapid import get_vapid_keypair
+
+
+class GenerateValidKeys(BaseCommand):
+    def handle(self, *args, **options):
+        pubkey, privkey = get_vapid_keypair()
+        self.stdout.write(f"Public key: {pubkey}")
+        self.stdout.write(f"Private key: {privkey}")
diff --git a/webpush/tests.py b/webpush/tests.py
deleted file mode 100644
index 7ce503c..0000000
--- a/webpush/tests.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from django.test import TestCase
-
-# Create your tests here.
diff --git a/webpush/tests/__init__.py b/webpush/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/webpush/tests/test_vapid.py b/webpush/tests/test_vapid.py
new file mode 100644
index 0000000..e5f4b38
--- /dev/null
+++ b/webpush/tests/test_vapid.py
@@ -0,0 +1,25 @@
+import unittest
+
+from cryptography.hazmat.primitives import serialization
+from py_vapid import Vapid, b64urlencode
+from webpush.vapid import get_vapid_keypair
+
+
+class TestGetVapidKeypair(unittest.TestCase):
+    def test_keypair_is_consistent(self):
+        # GIVEN a generated pair of public key and private key
+        pubkey, privkey = get_vapid_keypair()
+
+        # WHEN instantiating a new Vapid keypair from the private key
+        newkey = Vapid.from_raw(privkey.encode())
+
+        # THEN the public key export is the same as the one generated
+        self.assertEqual(
+            b64urlencode(
+                newkey.public_key.public_bytes(
+                    serialization.Encoding.X962,
+                    serialization.PublicFormat.UncompressedPoint,
+                )
+            ),
+            pubkey,
+        )
diff --git a/webpush/vapid.py b/webpush/vapid.py
new file mode 100644
index 0000000..ed0115b
--- /dev/null
+++ b/webpush/vapid.py
@@ -0,0 +1,16 @@
+from cryptography.hazmat.primitives import serialization
+from py_vapid import Vapid
+from py_vapid.utils import num_to_bytes, b64urlencode
+
+
+def get_vapid_keypair() -> tuple[str, str]:
+    vapid = Vapid()
+    vapid.generate_keys()
+
+    pubkey_export = vapid.public_key.public_bytes(
+        serialization.Encoding.X962,
+        serialization.PublicFormat.UncompressedPoint,
+    )
+    privkey_export = num_to_bytes(vapid.private_key.private_numbers().private_value, 32)
+
+    return b64urlencode(pubkey_export), b64urlencode(privkey_export)