From 281f9e8cb2b07942dc103eb2fe4dd45455beb1f4 Mon Sep 17 00:00:00 2001
From: JJTech0130 <jjtech@jjtech.dev>
Date: Sun, 23 Jul 2023 15:47:35 -0400
Subject: [PATCH] integrate it into the demo

---
 demo.py                      |  5 ++++-
 emulated/jelly.py            |  3 +++
 emulated/{test.py => nac.py} | 27 ++++++++++++++-------------
 3 files changed, 21 insertions(+), 14 deletions(-)
 rename emulated/{test.py => nac.py} (97%)

diff --git a/demo.py b/demo.py
index f477d3c..036b4ba 100644
--- a/demo.py
+++ b/demo.py
@@ -79,7 +79,10 @@ if CONFIG.get("id", {}).get("cert") is not None:
     id_keypair = ids._helpers.KeyPair(CONFIG["id"]["key"], CONFIG["id"]["cert"])
     user.restore_identity(id_keypair)
 else:
-    vd = input_multiline("Enter validation data: ")
+    #vd = input_multiline("Enter validation data: ")
+    import emulated.nac
+    vd = emulated.nac.generate_validation_data()
+    vd = b64encode(vd).decode()
     user.register(vd)
 
 print(user.lookup(["mailto:textgpt@icloud.com"]))
diff --git a/emulated/jelly.py b/emulated/jelly.py
index 383f391..5c58147 100644
--- a/emulated/jelly.py
+++ b/emulated/jelly.py
@@ -2,6 +2,9 @@ from io import BytesIO
 import unicorn
 from . import mparser as macholibre
 
+print = lambda *args, **kwargs: None
+
+
 STOP_ADDRESS = 0x00900000 # Used as a return address when calling functions
 
 ARG_REGISTERS = [
diff --git a/emulated/test.py b/emulated/nac.py
similarity index 97%
rename from emulated/test.py
rename to emulated/nac.py
index ebc8d9b..1592b6b 100644
--- a/emulated/test.py
+++ b/emulated/nac.py
@@ -350,11 +350,12 @@ def arc4random(j: Jelly) -> int:
     return random.randint(0, 0xFFFFFFFF)
     #return 0
 
-def main():
+def load_nac() -> Jelly:
     binary = load_binary()
     binary = get_x64_slice(binary)
     # Create a Jelly object from the binary
     j = Jelly(binary)
+
     hooks = {
         "_malloc": malloc,
         "___stack_chk_guard": lambda: 0,
@@ -395,22 +396,22 @@ def main():
         "_arc4random": arc4random
     }
     j.setup(hooks)
-    #j.uc.hook_add(unicorn.UC_HOOK_CODE, hook_code)
 
-    from base64 import b64encode
-    cert = get_cert()
-    val_ctx, req = nac_init(j,cert)
-    print(f"Validation Context: {hex(val_ctx)}")
-    print(f"Request: {b64encode(req).decode()}")
+    return j
 
+def generate_validation_data() -> bytes:
+    j = load_nac()
+    val_ctx, req = nac_init(j,get_cert())
     session_info = get_session_info(req)
-    print(f"Session Info: {b64encode(session_info).decode()}")
-
     nac_submit(j, val_ctx, session_info)
-
     val_data = nac_generate(j, val_ctx)
-
-    print(f"Validation Data: {b64encode(val_data).decode()}")
+    return bytes(val_data)
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    from base64 import b64encode
+    val_data = generate_validation_data()
+    print(f"Validation Data: {b64encode(val_data).decode()}")
+    #main()
+else:
+    # lazy hack: Disable print so that it's clean when not debugging
+    print = lambda *args, **kwargs: None
\ No newline at end of file