more refactoring

ids/__init__.py

@@ -1,10 +1,9 @@
 from base64 import b64encode
 import apns
-from . import profile
-from . import _helpers
-#from .profile import _get_auth_cert, _get_auth_token, _get_handles
+from . import profile, _helpers, identity
 
 class IDSUser:
+    # Sets self.user_id and self._auth_token
     def _authenticate_for_token(
         self, username: str, password: str, factor_callback: callable = None
     ):
@@ -12,6 +11,7 @@ class IDSUser:
             username, password, factor_callback
         )
 
+    # Sets self._auth_keypair using self.user_id and self._auth_token
     def _authenticate_for_cert(self):
         self._auth_keypair = profile._get_auth_cert(self.user_id, self._auth_token)
 
@@ -19,19 +19,39 @@ class IDSUser:
     def __init__(
         self,
         push_connection: apns.APNSConnection,
-        username: str,
-        password: str,
-        factor_callback: callable = None,
     ):
         self.push_connection = push_connection
+        self._push_keypair = _helpers.KeyPair(self.push_connection.private_key, self.push_connection.cert)
+
+    def __str__(self):
+        return f"IDSUser(user_id={self.user_id}, handles={self.handles}, push_token={b64encode(self.push_connection.token).decode()})"
+    
+    # Authenticates with a username and password, to create a brand new authentication keypair
+    def authenticate(self, username: str, password: str, factor_callback: callable = None):
         self._authenticate_for_token(username, password, factor_callback)
         self._authenticate_for_cert()
         self.handles = profile._get_handles(
             b64encode(self.push_connection.token),
             self.user_id,
             self._auth_keypair,
-            _helpers.KeyPair(self.push_connection.private_key, self.push_connection.cert),
+            self._push_keypair,
         )
 
-    def __str__(self):
-        return f"IDSUser(user_id={self.user_id}, handles={self.handles}, push_token={b64encode(self.push_connection.token).decode()})"
+    # Uses an existing authentication keypair
+    def restore_authentication(self, auth_keypair: _helpers.KeyPair, user_id: str, handles: dict):
+        self._auth_keypair = auth_keypair
+        self.user_id = user_id
+        self.handles = handles
+        
+    
+    # This is a separate call so that the user can make sure the first part succeeds before asking for validation data
+    def register(self, validation_data: str):
+        resp = identity.register_request(
+            b64encode(self.push_connection.token),
+            self.handles,
+            self.user_id,
+            self._auth_keypair,
+            self._push_keypair,
+            validation_data
+        )
+        print(resp)

ids/identity.py

@@ -0,0 +1,54 @@
+from ._helpers import KeyPair, PROTOCOL_VERSION, USER_AGENT
+from base64 import b64decode
+import plistlib
+import requests
+from .signing import add_auth_signature
+
+def register_request(
+    push_token, handles, uid, auth_key: KeyPair, push_key: KeyPair, validation_data
+):
+    uris = [{"uri": handle} for handle in handles]
+
+    body = {
+        "hardware-version": "MacBookPro18,3",
+        "language": "en-US",
+        "os-version": "macOS,13.2.1,22D68",
+        "software-version": "22D68",
+        "services": [
+            {
+                "capabilities": [{"flags": 1, "name": "Messenger", "version": 1}],
+                "service": "com.apple.madrid",
+                "users": [
+                    {
+                        # TODO: Pass ALL URIs from get handles
+                        "uris": uris,
+                        "user-id": uid,
+                    }
+                ],
+            }
+        ],
+        "validation-data": b64decode(validation_data),
+    }
+
+    body = plistlib.dumps(body)
+
+    headers = {
+        "x-protocol-version": PROTOCOL_VERSION,
+        "x-auth-user-id-0": uid,
+    }
+    add_auth_signature(
+        headers, body, "id-register", auth_key, push_key, push_token, 0
+    )
+
+    r = requests.post(
+        "https://identity.ess.apple.com/WebObjects/TDIdentityService.woa/wa/register",
+        headers=headers,
+        data=body,
+        verify=False,
+    )
+    r = plistlib.loads(r.content)
+    print(f'Response code: {r["status"]}')
+    if "status" in r and r["status"] == 6004:
+        raise Exception("Validation data expired!")
+    # TODO: Do validation of nested statuses
+    return r

newdemo.py

@@ -1,13 +1,59 @@
 import ids
 import apns
 from getpass import getpass
+import json
+from base64 import b64encode
+
+def input_multiline(prompt):
+    print(prompt)
+    lines = []
+    while True:
+        line = input()
+        if line == "":
+            break
+        lines.append(line)
+    return "\n".join(lines)
+
+# Try and load config.json
+try:
+    with open("config.json", "r") as f:
+        
+        CONFIG = json.load(f)
+except FileNotFoundError:
+    CONFIG = {}
 
 
-conn = apns.APNSConnection()
-conn.connect()
+conn = apns.APNSConnection(CONFIG.get("push", {}).get("key"), CONFIG.get("push", {}).get("cert"))
+conn.connect(CONFIG.get("push", {}).get("token"))
 
+user = ids.IDSUser(conn)
+
+if CONFIG.get("auth", {}).get("cert") is not None:
+    auth_keypair = ids._helpers.KeyPair(CONFIG["auth"]["key"], CONFIG["auth"]["cert"])
+    user_id = CONFIG["auth"]["user_id"]
+    handles = CONFIG["auth"]["handles"]
+    user.restore_authentication(auth_keypair, user_id, handles)
+else:
     username = input("Username: ")
     password = getpass("Password: ")
-user = ids.IDSUser(conn, username, password)
 
-print(user.handles)
+    user.authenticate(username, password)
+
+vd = input_multiline("Enter validation data: ")
+user.register(vd)
+
+# Write config.json
+CONFIG["auth"] = {
+    "key": user._auth_keypair.key,
+    "cert": user._auth_keypair.cert,
+    "user_id": user.user_id,
+    "handles": user.handles,
+}
+CONFIG["push"] = {
+    "token": b64encode(user.push_connection.token).decode(),
+    "key": user.push_connection.private_key,
+    "cert": user.push_connection.cert,
+}
+
+with open("config.json", "w") as f:
+    json.dump(CONFIG, f, indent=4)