From 3ff42f0cd490d9a028bf436332b71ee3fb002116 Mon Sep 17 00:00:00 2001
From: Ryan Barrett <git@ryanb.org>
Date: Thu, 31 Aug 2023 13:49:45 -0700
Subject: [PATCH] User.p256_key = k256_key, for ATProto

---
 models.py            | 21 +++++++++++++--------
 tests/test_models.py | 10 ++++++----
 tests/testutil.py    |  9 +++++----
 3 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/models.py b/models.py
index 9e2c9e5..619da5e 100644
--- a/models.py
+++ b/models.py
@@ -6,7 +6,9 @@ import logging
 import random
 from urllib.parse import quote, urlparse
 
-from Crypto.PublicKey import ECC, RSA
+import arroba.util
+from Crypto.PublicKey import RSA
+from cryptography.hazmat.primitives import serialization
 import dag_json
 from flask import g, request
 from google.cloud import ndb
@@ -78,15 +80,15 @@ class User(StringIdModel, metaclass=ProtocolUserMeta):
         section 5.1 of the Magic Signatures spec
       https://tools.ietf.org/html/draft-cavage-http-signatures-12
 
-    * P-256 keypair for AT Protocol's signing key
-      property: p256_key, PEM encoded
+    * K-256 keypair for AT Protocol's signing key
+      property: k256_key, PEM encoded
       https://atproto.com/guides/overview#account-portability
     """
     obj_key = ndb.KeyProperty(kind='Object')  # user profile
     mod = ndb.StringProperty()
     public_exponent = ndb.StringProperty()
     private_exponent = ndb.StringProperty()
-    p256_key = ndb.StringProperty()
+    k256_key = ndb.BlobProperty()
     use_instead = ndb.KeyProperty()
 
     # whether this user signed up or otherwise explicitly, deliberately
@@ -159,10 +161,13 @@ class User(StringIdModel, metaclass=ProtocolUserMeta):
                     'private_exponent': long_to_base64(key.d),
             })
 
-        if cls.LABEL != 'atprotocol':
-            key = ECC.generate(
-                curve='P-256', randfunc=random.randbytes if DEBUG else None)
-            kwargs['p256_key'] = key.export_key(format='PEM')
+        if cls.LABEL != 'atproto':
+            privkey = arroba.util.new_key()
+            kwargs['k256_key'] = privkey.private_bytes(
+                encoding=serialization.Encoding.PEM,
+                format=serialization.PrivateFormat.PKCS8,
+                encryption_algorithm=serialization.NoEncryption(),
+            )
 
         user = cls(id=id, **kwargs)
         try:
diff --git a/tests/test_models.py b/tests/test_models.py
index e397473..5e74185 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -11,6 +11,8 @@ from oauth_dropins.webutil.testutil import NOW
 from .testutil import Fake, TestCase
 
 from atproto import ATProto
+from cryptography.hazmat.primitives.asymmetric import ec
+from cryptography.hazmat.primitives import serialization
 from models import Follower, Object, OBJECT_EXPIRE_AGE, Target, User
 import protocol
 from protocol import Protocol
@@ -32,15 +34,15 @@ class UserTest(TestCase):
         assert user.mod
         assert user.public_exponent
         assert user.private_exponent
-        assert user.p256_key
+        assert user.k256_key
 
         # check that we can load the keys
         assert user.public_pem()
         assert user.private_pem()
 
-        p256_key = ECC.import_key(user.p256_key)
-        assert isinstance(p256_key, ECC.EccKey)
-        self.assertEqual('NIST P-256', p256_key.curve)
+        k256_key = serialization.load_pem_private_key(user.k256_key, password=None)
+        self.assertIsInstance(k256_key, ec.EllipticCurvePrivateKey)
+        self.assertIsInstance(k256_key.curve, ec.SECP256K1)
 
         # direct should get set even if the user exists
         same = Fake.get_or_create('a.b', direct=True)
diff --git a/tests/testutil.py b/tests/testutil.py
index 7e40a11..f08f01c 100644
--- a/tests/testutil.py
+++ b/tests/testutil.py
@@ -142,6 +142,7 @@ models.reset_protocol_properties()
 
 import app
 from activitypub import ActivityPub, CONNEG_HEADERS_AS2_HTML
+from atproto import ATProto
 import common
 from web import Web
 from flask_app import app, cache, init_globals
@@ -250,7 +251,7 @@ class TestCase(unittest.TestCase, testutil.Asserts):
                    mod=global_user.mod,
                    public_exponent=global_user.public_exponent,
                    private_exponent=global_user.private_exponent,
-                   p256_key=global_user.p256_key,
+                   k256_key=global_user.k256_key,
                    obj_key=obj_key,
                    **kwargs)
         user.put()
@@ -408,7 +409,7 @@ class TestCase(unittest.TestCase, testutil.Asserts):
             self.assert_equals(obj_as2, got.as2())
 
         # generated, computed, etc
-        ignore = ['created', 'mod', 'obj_key', 'p256_key', 'private_exponent',
+        ignore = ['created', 'mod', 'obj_key', 'k256_key', 'private_exponent',
                   'public_exponent', 'readable_id', 'updated']
         for prop in ignore:
             assert prop not in props
@@ -420,8 +421,8 @@ class TestCase(unittest.TestCase, testutil.Asserts):
             assert got.private_exponent
             assert got.public_exponent
 
-        # if cls != ATProto:
-        #     assert got.p256_key
+        if cls != ATProto:
+            assert got.k256_key
 
         return got