From 2cab8c4c9240ad00211db4f38fcfd7b4a0661886 Mon Sep 17 00:00:00 2001
From: Lil-Ran <lilran4869@qq.com>
Date: Thu, 27 Feb 2025 16:10:30 +0800
Subject: [PATCH] wip: header parsing

---
 helpers/shot.py |  2 +-
 pyc_code.h      |  3 +++
 pyc_module.cpp  | 69 +++++++++++++++++++++++++++++++++++++++++++++++++
 pyc_module.h    |  6 +++++
 pycdas.cpp      |  2 +-
 5 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/helpers/shot.py b/helpers/shot.py
index 3ed0a57..6bef35b 100644
--- a/helpers/shot.py
+++ b/helpers/shot.py
@@ -37,7 +37,7 @@ def decrypt_process(runtimes: dict[str, RuntimeInfo], sequences: list[tuple[str,
             with open(dest_path + '.1shot.seq', 'wb') as f:
                 f.write(b'\xa1' + runtime.runtime_aes_key)
                 f.write(b'\xa2' + runtime.mix_str_aes_nonce())
-                f.write(b'\xf0\xf0')
+                f.write(b'\xf0\xff')
                 f.write(data[:cipher_text_offset])
                 f.write(general_aes_ctr_decrypt(
                     data[cipher_text_offset:cipher_text_offset+cipher_text_length], runtime.runtime_aes_key, nonce))
diff --git a/pyc_code.h b/pyc_code.h
index e6b2ce9..d70fe98 100644
--- a/pyc_code.h
+++ b/pyc_code.h
@@ -35,6 +35,9 @@ public:
         CO_FUTURE_GENERATOR_STOP = 0x800000,                // 3.5 ->
         CO_FUTURE_ANNOTATIONS = 0x1000000,                  // 3.7 ->
         CO_NO_MONITORING_EVENTS = 0x2000000,                // 3.13 ->
+
+        // TODO: Shift things
+        CO_PYARMOR_OBFUSCATED = 0x20000000,
     };
 
     PycCode(int type = TYPE_CODE)
diff --git a/pyc_module.cpp b/pyc_module.cpp
index 98eecff..70744ad 100644
--- a/pyc_module.cpp
+++ b/pyc_module.cpp
@@ -251,6 +251,75 @@ void PycModule::loadFromMarshalledFile(const char* filename, int major, int mino
     m_code = LoadObject(&in, this).cast<PycCode>();
 }
 
+void PycModule::loadFromOneshotSequenceFile(const char *filename)
+{
+    PycFile in(filename);
+    if (!in.isOpen())
+    {
+        fprintf(stderr, "Error opening file %s\n", filename);
+        return;
+    }
+
+    bool oneshot_seq_header = true;
+    while (oneshot_seq_header)
+    {
+        int indicator = in.getByte();
+        switch (indicator)
+        {
+        case 0xA1:
+            in.getBuffer(16, this->pyarmor_aes_key);
+            break;
+        case 0xA2:
+            in.getBuffer(12, this->pyarmor_mix_str_aes_nonce);
+            break;
+        case 0xF0:
+            break;
+        case 0xFF:
+            oneshot_seq_header = false;
+            break;
+        default:
+            fprintf(stderr, "Unknown 1-shot sequence indicator %02X\n", indicator);
+            break;
+        }
+    }
+
+    // Write only. Some fields unknown to us or not needed for decryption are discarded.
+    char discard_buffer[64];
+
+    char pyarmor_header[64];
+    in.getBuffer(64, pyarmor_header);
+    this->m_maj = pyarmor_header[9];
+    this->m_min = pyarmor_header[10];
+    this->m_unicode = (m_maj >= 3);
+
+    unsigned int remain_header_length = *(unsigned int *)(pyarmor_header + 28) - 64;
+    while (remain_header_length)
+    {
+        unsigned int discard_length = (remain_header_length > 64) ? 64 : remain_header_length;
+        in.getBuffer(discard_length, discard_buffer);
+        remain_header_length -= discard_length;
+    }
+
+    // For 1-shot sequence, the following part has been decrypted once.
+    unsigned int code_object_offset = in.get32();
+    unsigned int co_code_aes_nonce_xor_key_procedure_length = in.get32();
+    this->pyarmor_co_code_aes_nonce_xor_enabled = (co_code_aes_nonce_xor_key_procedure_length > 0);
+    unsigned int remain_second_part_length = code_object_offset - 8;
+    while (remain_second_part_length)
+    {
+        unsigned int discard_length = (remain_second_part_length > 64) ? 64 : remain_second_part_length;
+        in.getBuffer(discard_length, discard_buffer);
+        remain_second_part_length -= discard_length;
+    }
+
+    if (this->pyarmor_co_code_aes_nonce_xor_enabled)
+    {
+        // TODO: Implement the decryption procedure.
+    }
+
+    m_code = LoadObject(&in, this).cast<PycCode>();
+}
+
 PycRef<PycString> PycModule::getIntern(int ref) const
 {
     if (ref < 0 || (size_t)ref >= m_interns.size())
diff --git a/pyc_module.h b/pyc_module.h
index 695f4ba..36248d4 100644
--- a/pyc_module.h
+++ b/pyc_module.h
@@ -46,6 +46,7 @@ public:
 
     void loadFromFile(const char* filename);
     void loadFromMarshalledFile(const char *filename, int major, int minor);
+    void loadFromOneshotSequenceFile(const char* filename);
     bool isValid() const { return (m_maj >= 0) && (m_min >= 0); }
 
     int majorVer() const { return m_maj; }
@@ -87,6 +88,11 @@ private:
     int m_maj, m_min;
     bool m_unicode;
 
+    char pyarmor_aes_key[16];
+    char pyarmor_mix_str_aes_nonce[12];
+    bool pyarmor_co_code_aes_nonce_xor_enabled;
+    char pyarmor_co_code_aes_nonce_xor_key[12];
+
     PycRef<PycCode> m_code;
     std::vector<PycRef<PycString>> m_interns;
     std::vector<PycRef<PycObject>> m_refs;
diff --git a/pycdas.cpp b/pycdas.cpp
index b73410f..2b4a6fe 100644
--- a/pycdas.cpp
+++ b/pycdas.cpp
@@ -23,7 +23,7 @@ static const char* flag_names[] = {
     "CO_FUTURE_PRINT_FUNCTION", "CO_FUTURE_UNICODE_LITERALS", "CO_FUTURE_BARRY_AS_BDFL",
             "CO_FUTURE_GENERATOR_STOP",
     "CO_FUTURE_ANNOTATIONS", "CO_NO_MONITORING_EVENTS", "<0x4000000>", "<0x8000000>",
-    "<0x10000000>", "<0x20000000>", "<0x40000000>", "<0x80000000>"
+    "<0x10000000>", "CO_PYARMOR_OBFUSCATED", "<0x40000000>", "<0x80000000>"
 };
 
 static void print_coflags(unsigned long flags, std::ostream& pyc_output)