working demo!

demo.py

@@ -6,90 +6,121 @@ import json
 # Open config
 try:
     with open("config.json", "r") as f:
-        config = json.load(f)
+        CONFIG = json.load(f)
 except FileNotFoundError:
-    config = {}
+    CONFIG = {}
 
-# If no username is set, prompt for it
+def refresh_token():
-if "username" not in config:
+    # If no username is set, prompt for it
-    config["username"] = input("Enter iCloud username: ")
+    if "username" not in CONFIG:
-# If no password is set, prompt for it
+        CONFIG["username"] = input("Enter iCloud username: ")
-if "password" not in config:
+    # If no password is set, prompt for it
-    config["password"] = getpass.getpass("Enter iCloud password: ")
+    if "password" not in CONFIG:
-# If grandslam authentication is not set, prompt for it
+        CONFIG["password"] = getpass.getpass("Enter iCloud password: ")
-if "use_gsa" not in config:
+    # If grandslam authentication is not set, prompt for it
-    config["use_gsa"] = input("Use grandslam authentication? [y/N] ").lower() == "y"
+    if "use_gsa" not in CONFIG:
+        CONFIG["use_gsa"] = input("Use grandslam authentication? [y/N] ").lower() == "y"
 
-def factor_gen():
+    def factor_gen():
-    return input("Enter iCloud 2FA code: ")
+        return input("Enter iCloud 2FA code: ")
 
-user_id, token = ids._get_auth_token(
+    CONFIG["user_id"], CONFIG["token"] = ids._get_auth_token(
-    config["username"], config["password"], config["use_gsa"], factor_gen=factor_gen
+        CONFIG["username"], CONFIG["password"], CONFIG["use_gsa"], factor_gen=factor_gen
-)
+    )
 
-config["user_id"] = user_id
+def refresh_cert():
-config["token"] = token
+    CONFIG["key"], CONFIG["auth_cert"] = ids._get_auth_cert(
+        CONFIG["user_id"], CONFIG["token"]
+    )
 
-key, cert = ids._get_auth_cert(user_id, token)
+def create_connection(): 
+    conn = apns.APNSConnection()
+    token = conn.connect()
+    conn.filter(['com.apple.madrid'])
+    CONFIG['push'] = {
+        'token': b64encode(token).decode(),
+        'cert': conn.cert,
+        'key': conn.private_key
+    }
+    return conn
 
-config["key"] = key
+def restore_connection():
-config["cert"] = cert
+    conn = apns.APNSConnection(CONFIG['push']['key'], CONFIG['push']['cert'])
+    conn.connect(True, b64decode(CONFIG['push']['token']))
+    conn.filter(['com.apple.madrid'])
+    return conn
 
-conn1 = apns.APNSConnection()
+def refresh_madrid_cert():
-conn1.connect()
+    info = {
+        "uri": "mailto:" + CONFIG["username"],
+        "user_id": CONFIG['user_id'],
+    }
 
-conn1.filter(["com.apple.madrid"])
-
-info = {
-    "uri": "mailto:" + config["username"],
-    "user_id": user_id,
-}
-
-resp = None
-try:
-    if "validation_data" in config:
-        resp = ids._register_request(
-            b64encode(conn1.token),
-            info,
-            cert,
-            key,
-            conn1.cert,
-            conn1.private_key,
-            config["validation_data"],
-        )
-except Exception as e:
-    print(e)
     resp = None
+    try:
+        if "validation_data" in CONFIG:
+            resp = ids._register_request(
+                CONFIG['push']['token'],
+                info,
+                CONFIG['auth_cert'],
+                CONFIG['key'],
+                CONFIG['push']['cert'],
+                CONFIG['push']['key'],
+                CONFIG["validation_data"],
+            )
+    except Exception as e:
+        print(e)
+        resp = None
 
-if resp is None:
+    if resp is None:
-    print(
+        print(
-        "Note: Validation data can be obtained from @JJTech, or intercepted using a HTTP proxy."
+            "Note: Validation data can be obtained from @JJTech, or intercepted using a HTTP proxy."
-    )
+        )
-    validation_data = (
+        validation_data = (
-        input_multiline("Enter validation data: ")
+            input_multiline("Enter validation data: ")
-        .replace("\n", "")
+            .replace("\n", "")
-        .replace(" ", "")
+            .replace(" ", "")
-    )
+        )
-    resp = ids._register_request(
+        resp = ids._register_request(
-        b64encode(conn1.token),
+            CONFIG['push']['token'],
-        info,
+            info,
-        cert,
+            CONFIG['auth_cert'],
-        key,
+            CONFIG['key'],
-        conn1.cert,
+            CONFIG['push']['cert'],
-        conn1.private_key,
+            CONFIG['push']['key'],
-        validation_data,
+            validation_data,
-    )
+        )
-    config["validation_data"] = validation_data
+        CONFIG["validation_data"] = validation_data
 
-madrid_cert = x509.load_der_x509_certificate(
+    madrid_cert = x509.load_der_x509_certificate(
-    resp["services"][0]["users"][0]["cert"]
+        resp["services"][0]["users"][0]["cert"]
-)
+    )
-madrid_cert = (
+    madrid_cert = (
-    madrid_cert.public_bytes(serialization.Encoding.PEM).decode("utf-8").strip()
+        madrid_cert.public_bytes(serialization.Encoding.PEM).decode("utf-8").strip()
-)
+    )
 
-config["madrid_cert"] = madrid_cert
+    CONFIG["madrid_cert"] = madrid_cert
+
+
+if not 'madrid_cert' in CONFIG:
+    print("No madrid cert")
+    if not 'key' in CONFIG:
+        print("No auth cert")
+        if not 'token' in CONFIG:
+            print("No auth token")
+            refresh_token()
+        refresh_cert()
+    if not 'push' in CONFIG:
+        print("No push conn")
+        conn = create_connection()
+    else:
+        print("restoring push conn")
+        conn = restore_connection()
+    refresh_madrid_cert()
+    print("Got new madrid cert")
+
+print("Done")
 
 # Save config
 with open("config.json", "w") as f:
-    json.dump(config, f, indent=4)
+    json.dump(CONFIG, f, indent=4)