Add new skip option

CHANGELOG.md

@@ -4,6 +4,13 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [2.1.0] - Not release
+
+### Added
+
+- Added private key function.
+- Option to Skip automatic detection of private function.
+
 ## [2.0.9] - 2024-09-25
 
 ### Added

README.md

@@ -72,6 +72,7 @@ Cdm options:
                         Output directory path for extracted data.
   -f <file>, --functions <file>
                         Path to Ghidra XML functions file.
+  -s, --skip            Skip auto-detect of private function.
 
 ```
 

keydive/__main__.py

@@ -84,6 +84,7 @@ def main() -> None:
     opt_cdm.add_argument('-w', '--wvd', required=False, action='store_true', help='Generate a pywidevine WVD device file.')
     opt_cdm.add_argument('-o', '--output', required=False, type=Path, default=Path('device'), metavar='<dir>', help='Output directory path for extracted data.')
     opt_cdm.add_argument('-f', '--functions', required=False, type=Path, metavar='<file>', help='Path to Ghidra XML functions file.')
+    opt_cdm.add_argument('-s', '--skip', required=False, action='store_true', help='Skip auto-detect of private function.')
     args = parser.parse_args()
 
     if args.version:
@@ -107,7 +108,7 @@ def main() -> None:
             cdm.set_challenge(data=args.challenge)
 
         # Initialize Core instance for interacting with the device
-        core = Core(cdm=cdm, device=args.device, functions=args.functions)
+        core = Core(cdm=cdm, device=args.device, functions=args.functions, skip=args.skip)
 
         # Process watcher loop
         logger.info('Watcher delay: %ss' % args.delay)

keydive/core.py

@@ -20,7 +20,7 @@ class Core:
     Core class for handling DRM operations and device interactions.
     """
 
-    def __init__(self, cdm: Cdm, device: str = None, functions: Path = None):
+    def __init__(self, cdm: Cdm, device: str = None, functions: Path = None, skip: bool = False):
         """
         Initializes a Core instance.
 
@@ -28,10 +28,12 @@ class Core:
             cdm (Cdm): Instance of Cdm for managing DRM related operations.
             device (str, optional): ID of the Android device to connect to via ADB. Defaults to None (uses USB device).
             functions (Path, optional): Path to Ghidra XML functions file for symbol extraction. Defaults to None.
+            skip (bool, optional): Flag to determine whether to skip predefined functions (e.g., OEM_CRYPTO_API).
         """
         self.logger = logging.getLogger(self.__class__.__name__)
         self.running = True
         self.cdm = cdm
+        self.skip = skip
 
         # Select device based on provided ID or default to the first USB device.
         self.device: Device = frida.get_device(id=device, timeout=5) if device else frida.get_usb_device(timeout=5)
@@ -61,7 +63,8 @@ class Core:
         replacements = {
             '${OEM_CRYPTO_API}': json.dumps(list(OEM_CRYPTO_API)),
             '${NATIVE_C_API}': json.dumps(list(NATIVE_C_API)),
-            '${SYMBOLS}': json.dumps(symbols)
+            '${SYMBOLS}': json.dumps(symbols),
+            '${SKIP}': str(self.skip)
         }
 
         for placeholder, value in replacements.items():
@@ -69,8 +72,7 @@ class Core:
 
         return content
 
-    @staticmethod
-    def __prepare_symbols(path: Path) -> list:
+    def __prepare_symbols(self, path: Path) -> list:
         """
         Parses the provided XML functions file to select relevant functions.
 
@@ -95,7 +97,7 @@ class Core:
             functions = program['FUNCTIONS']['FUNCTION']
 
             # Find a target function from a predefined list
-            target = next((f['@NAME'] for f in functions if f['@NAME'] in OEM_CRYPTO_API), None)
+            target = None if self.skip else next((f['@NAME'] for f in functions if f['@NAME'] in OEM_CRYPTO_API), None)
 
             # Extract relevant functions
             selected = {}
@@ -106,7 +108,7 @@ class Core:
                 # Add function if it matches specific criteria
                 if name not in selected and (
                         name == target
-                        or any(keyword in name for keyword in CDM_FUNCTION_API)
+                        or any(None if self.skip else keyword in name for keyword in CDM_FUNCTION_API)
                         or (not target and re.match(r'^[a-z]+$', name) and args >= 6)
                 ):
                     selected[name] = {

keydive/keydive.js

@@ -8,6 +8,7 @@
 const OEM_CRYPTO_API = JSON.parse('${OEM_CRYPTO_API}');
 const NATIVE_C_API = JSON.parse('${NATIVE_C_API}');
 const SYMBOLS = JSON.parse('${SYMBOLS}');
+const SKIP = '${SKIP}' === 'True';
 
 
 // Logging levels to synchronize with Python's logging module.
@@ -281,7 +282,7 @@ const hookLibrary = (name) => {
     }
 
     functions = functions.filter(f => !NATIVE_C_API.includes(f.name));
-    const targets = functions.filter(f => OEM_CRYPTO_API.includes(f.name)).map(f => f.name);
+    const targets = SKIP ? [] : functions.filter(f => OEM_CRYPTO_API.includes(f.name)).map(f => f.name);
     const hooked = [];
 
     functions.forEach(func => {