Add tests

dhinakg · Jun 22, 2024 · 883718e · 883718e
1 parent 5544f44
commit 883718e
Show file tree

Hide file tree

Showing 5 changed files with 189 additions and 50 deletions.
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -0,0 +1,86 @@
+name: Build
+
+on:
+  push:
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  build:
+    strategy:
+      matrix:
+        os: [macos-12, macos-13, macos-14]
+    runs-on: ${{ matrix.os }}
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Build
+        run: make
+      - name: Test binary
+        run: |
+          ./aastuff
+          ./aastuff_standalone
+      - name: Upload artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: macOS artifacts for ${{ matrix.os }}
+          path: aastuff*
+
+  test-get-key:
+    strategy:
+      matrix:
+        python-version: [3.10, 3.11, 3.12]
+    runs-on: [ubuntu-latest, macos-latest, windows-latest]
+
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: "pip"
+      - name: Cache test files
+        id: cache
+        uses: actions/cache@v4
+        with:
+          key: test-get-key-files
+          path: |
+            tests
+      - name: Download test files
+        if: ${{ steps.cache.outputs.cache-hit != 'true' }}
+        run: |
+          mkdir tests tests/macOS_15_beta_1_OTA tests/iOS_18_beta_1_IPSW/
+          curl -L "https://updates.cdn-apple.com/2024SummerSeed/mobileassets/052-49061/CA7135A8-BAF6-4890-887C-35FB30C154D5/com_apple_MobileAsset_MacSoftwareUpdate/e2de87f20576b2bdc021d36f74a2f836cf42afe576178388dfd0cde875f4f979.aea" -o tests/macOS_15_beta_1_OTA/encrypted.aea
+          echo "$MACOS_OTA_TEST_KEY" > tests/macOS_15_beta_1_OTA/expected.txt
+          curl -L "https://updates.cdn-apple.com/2024SummerSeed/fullrestores/052-34764/D5D3D10C-E557-4A46-8EBD-290411A228AA/iPhone16,2_18.0_22A5282m_Restore.ipsw" -o tests/iPhone_15PM_18.0_22A5282m.ipsw
+          unzip -p tests/iPhone_15PM_18.0_22A5282m.ipsw 090-29713-049.dmg.aea > tests/iOS_18_beta_1_IPSW/encrypted.aea
+          rm tests/iPhone_15PM_18.0_22A5282m.ipsw
+          echo "$IOS_IPSW_TEST_KEY" > tests/iOS_18_beta_1_IPSW/expected.txt
+      - name: Run tests
+        run: ./test_get_key.sh
+
+  test-aastuff:
+    runs-on: macos-latest
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Cache test files
+        id: cache
+        uses: actions/cache@v4
+        with:
+          key: test-aastuff-files
+          path: |
+            tests
+      - name: Download test files
+        if: ${{ steps.cache.outputs.cache-hit != 'true' }}
+        run: |
+          mkdir tests tests/small tests/large
+          # This file uses a compressed inner layer 
+          curl -L "https://updates.cdn-apple.com/2024/Iris/mobileassets/003-49672/A1233F60-3D17-491B-803A-DB26E20695AE/com_apple_MobileAsset_UAF_Siri_Understanding/6FF3BAF0-FBEF-4C01-BB0E-30CD61DAFCC4.aar" -o tests/small/encrypted.aea
+          echo "$SMALL_TEST_KEY" > tests/small/key.txt
+          # This file uses a raw inner layer
+          curl -L "https://updates.cdn-apple.com/2024SummerSeed/mobileassets/052-49061/CA7135A8-BAF6-4890-887C-35FB30C154D5/com_apple_MobileAsset_MacSoftwareUpdate/e2de87f20576b2bdc021d36f74a2f836cf42afe576178388dfd0cde875f4f979.aea" -o tests/large/encrypted.aea
+          echo "$LARGE_TEST_KEY" > tests/large/key.txt
+      - name: Build
+        run: make
+      - name: Run tests
+        run: ./test_extract.sh
diff --git a/get_key.py b/get_key.py
@@ -8,7 +8,8 @@
 import base64
 import json
 from pathlib import Path
-from sys import argv
+import argparse
+import sys
 
 import requests
 from pyhpke import AEADId, CipherSuite, KDFId, KEMId, KEMKey
@@ -17,72 +18,78 @@
 
 suite = CipherSuite.new(KEMId.DHKEM_P256_HKDF_SHA256, KDFId.HKDF_SHA256, AEADId.AES256_GCM)
 
-if len(argv) < 2:
-    print("Usage: get_key.py <aea>")
+
+def error(msg):
+    print(msg, file=sys.stderr)
     exit(1)
 
 
-aea_path = Path(argv[1])
-fields = {}
-with aea_path.open("rb") as f:
-    header = f.read(12)
-    if len(header) != 12:
-        print(f"Expected 12 bytes, got {len(header)}")
-        exit(1)
+def main(aea_path: Path, verbose: bool = False):
+    fields = {}
+    with aea_path.open("rb") as f:
+        header = f.read(12)
+        if len(header) != 12:
+            error(f"Expected 12 bytes, got {len(header)}")
 
-    magic = header[:4]
-    if magic != b"AEA1":
-        print(f"Invalid magic: {magic.hex()}")
-        exit(1)
+        magic = header[:4]
+        if magic != b"AEA1":
+            error(f"Invalid magic: {magic.hex()}")
 
-    profile = int.from_bytes(header[4:7], "little")
-    if profile != AEA_PROFILE__HKDF_SHA256_AESCTR_HMAC__SYMMETRIC__NONE:
-        print(f"Invalid AEA profile: {profile}")
-        exit(1)
+        profile = int.from_bytes(header[4:7], "little")
+        if profile != AEA_PROFILE__HKDF_SHA256_AESCTR_HMAC__SYMMETRIC__NONE:
+            error(f"Invalid AEA profile: {profile}")
 
-    auth_data_blob_size = int.from_bytes(header[8:12], "little")
+        auth_data_blob_size = int.from_bytes(header[8:12], "little")
 
-    if auth_data_blob_size == 0:
-        print("No auth data blob")
-        exit(1)
+        if auth_data_blob_size == 0:
+            error("No auth data blob")
 
-    auth_data_blob = f.read(auth_data_blob_size)
-    if len(auth_data_blob) != auth_data_blob_size:
-        print(f"Expected {auth_data_blob_size} bytes, got {len(auth_data_blob)}")
-        exit(1)
+        auth_data_blob = f.read(auth_data_blob_size)
+        if len(auth_data_blob) != auth_data_blob_size:
+            error(f"Expected {auth_data_blob_size} bytes, got {len(auth_data_blob)}")
 
-    assert auth_data_blob[:4]
+        assert auth_data_blob[:4]
 
-    while len(auth_data_blob) > 0:
-        field_size = int.from_bytes(auth_data_blob[:4], "little")
-        field_blob = auth_data_blob[:field_size]
+        while len(auth_data_blob) > 0:
+            field_size = int.from_bytes(auth_data_blob[:4], "little")
+            field_blob = auth_data_blob[:field_size]
 
-        key_end = field_blob.index(b"\x00", 4)
-        key = field_blob[4:key_end].decode()
-        value = field_blob[key_end + 1 :].decode()
-        fields[key] = value
+            key_end = field_blob.index(b"\x00", 4)
+            key = field_blob[4:key_end].decode()
+            value = field_blob[key_end + 1 :].decode()
+            fields[key] = value
 
-        auth_data_blob = auth_data_blob[field_size:]
+            auth_data_blob = auth_data_blob[field_size:]
 
+    if verbose:
+        print(fields, "\n", file=sys.stderr)
 
-print(fields, "\n")
+    if "com.apple.wkms.fcs-response" not in fields:
+        error("No fcs-response field found, is this from an OTA?")
 
-if "com.apple.wkms.fcs-response" not in fields:
-    print("No fcs-response field found, is this from an OTA?")
-    exit(1)
+    fcs_response = json.loads(fields["com.apple.wkms.fcs-response"])
+    enc_request = base64.b64decode(fcs_response["enc-request"])
+    wrapped_key = base64.b64decode(fcs_response["wrapped-key"])
+    url = fields["com.apple.wkms.fcs-key-url"]
+
+    r = requests.get(url, timeout=10)
+    r.raise_for_status()
 
-fcs_response = json.loads(fields["com.apple.wkms.fcs-response"])
-enc_request = base64.b64decode(fcs_response["enc-request"])
-wrapped_key = base64.b64decode(fcs_response["wrapped-key"])
-url = fields["com.apple.wkms.fcs-key-url"]
+    privkey = KEMKey.from_pem(r.text)
 
-r = requests.get(url, timeout=10)
-r.raise_for_status()
+    recipient = suite.create_recipient_context(enc_request, privkey)
+    pt = recipient.open(wrapped_key)
 
+    if verbose:
+        print(f"Key: {base64.b64encode(pt).decode()}")
+    else:
+        print(base64.b64encode(pt).decode())
 
-privkey = KEMKey.from_pem(r.text)
 
-recipient = suite.create_recipient_context(enc_request, privkey)
-pt = recipient.open(wrapped_key)
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Get the key for an AEA file")
+    parser.add_argument("path", help="Path to the AEA file")
+    parser.add_argument("-v", "--verbose", action="store_true", help="Show verbose output")
+    args = parser.parse_args()
 
-print(f"Key: {base64.b64encode(pt).decode()}")
+    main(Path(args.path), args.verbose)
diff --git a/src/aastuff.m b/src/aastuff.m
@@ -15,7 +15,7 @@ int main(int argc, char** argv) {
         if (argc < 3) {
             ERRLOG(@"Usage: %s <archive> <output directory> [key in base64]", argv[0]);
             ERRLOG(@"Key is required for encrypted archives");
-            return 1;
+            return argc == 1 ? 0 : 1;
         }
 
         NSString* archivePath = [NSString stringWithUTF8String:argv[1]];

diff --git a/test_extract.sh b/test_extract.sh
@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+
+set -e
+
+rm -rf tmp
+mkdir tmp
+
+for i in tests/*; do
+    echo "Testing $i"
+    mkdir -p "tmp/$i/a" "tmp/$i/b"
+    ./aastuff "tests/$i/encrypted.aar" "tmp/$i/a" "$(cat tests/"$i"/key.txt)"
+    ./aastuff_standalone "tests/$i/encrypted.aar" "tmp/$i/b" "$(cat tests/"$i"/key.txt)"
+    diff -r "tmp/$i/a" "tmp/$i/b" && echo "Test $i passed" || echo "Test $i failed"
+done
+
+rm -rf tmp
+echo Done
diff --git a/test_get_key.sh b/test_get_key.sh
@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+
+set -e
+
+abort() {
+    echo "$1"
+    exit 1
+}
+
+rm -rf tmp
+mkdir tmp
+
+for i in tests/*; do
+    echo "Testing $i"
+    mkdir -p "tmp/$i"
+    # Ensure expected key is valid first
+    aea decrypt -i "tests/$i/encrypted.aea" -o "tmp/decrypted" -key-value "base64:$(cat tests/"$i"/expected.txt)" || abort "Failed to decrypt with expected key"
+    # Get the key
+    python3 get_key.py "tests/$i/encrypted.aea" > "tmp/$1/actual.txt" || abort "Failed to get key"
+    # Ensure the key is correct
+    aea decrypt -i "tests/$i/encrypted.aea" -o "tmp/decrypted" -key-value "base64:$(cat tests/"$i"/actual.txt)" || abort "Failed to decrypt with actual key"
+    if diff -q "tmp/$i/actual.txt" "tests/$i/expected.txt"; then
+        echo "Warning: key does not match expected key, but decryption was successful"
+    fi
+    echo "Test $i passed"
+done
+
+rm -rf tmp
+echo Done