✨ feat(Python): Python hook

Add hooker to execute shell cmds in python directily as builtin functions
wychlw · Sep 8, 2024 · 2683ab3 · 2683ab3
1 parent 92d68ac
commit 2683ab3
Show file tree

Hide file tree

Showing 8 changed files with 79 additions and 40 deletions.
diff --git a/python/tester/__init__.py b/python/tester/__init__.py
@@ -3,4 +3,6 @@
 __doc__ = tester.__doc__
 if hasattr(tester, "__all__"):
     __all__ = tester.__all__
-
+
+from .hook import HookGlobals
+from .direct_script import DirectScript
diff --git a/python/tester/direct_script.py b/python/tester/direct_script.py
@@ -0,0 +1,20 @@
+from .hook import HookGlobals
+from . import tester
+
+class DirectScript(HookGlobals):
+    def hook_func(func, *args, **kwargs):
+        if globals().get(func) is not None:
+            return globals()[func](*args, **kwargs)
+        cmd = f"{func[1][1]} "
+        for arg in args:
+            cmd += f"{arg} "
+        for key, value in kwargs.items():
+            cmd += f"{key}={value} "
+        global __hook_exec__
+        __hook_exec__.script_run(cmd)
+
+    def __init__(self, func, exec):
+        super().__init__(func)
+        self.set_global("__hook_func__", DirectScript.hook_func)
+        self.set_global("__hook_exec__", exec)
+
diff --git a/python/tester/src/hook.py → python/tester/hook.py b/python/tester/src/hook.py → python/tester/hook.py
@@ -1,8 +1,9 @@
 from typing import Generator, Tuple
 from bytecode import *
+import inspect
 
 """
-对于我们想要 hook 的东西，其大概会长成这个样子：
+对于这里我们想要 hook 的东西，其大概会长成这个样子：
 ```python
 def func():
     ping("127.0.0.1")
@@ -25,7 +26,6 @@ def hook_func('name', *args, **kwargs):
 
 也就是说，我们需要去找这样一个 LOAD_GLOBAL 和 CALL 的组合，并且中间的所有 LOADxxx 和 arg 数量一致。而后，修改该段字节码，使得其实际调用的是我们的 hook_func。
 注意需要考虑多个 ARGS 和 KWARGS 的问题...
-
 """
 
 class HookGlobals:
@@ -62,26 +62,34 @@ def find_paired_call(self, code: Bytecode) -> Generator[Tuple[int, int], None, N
 
     def replace_global_call(self, func):
         codes = Bytecode.from_code(func.__code__)
-        for i in enumerate(codes):
-            print(i)
         for i, j in self.find_paired_call(codes):
-            print(i, j, ':')
-            for k in range(i, j + 1):
-                print(codes[k])
-            print()
+            orig_func = codes[i].arg
             replace_code = [
-                Instr('LOAD_GLOBAL', (True, '__hook_func__'))
+                Instr('LOAD_GLOBAL', (True, '__hook_func__')),
+                Instr('LOAD_CONST', (True, orig_func)),
             ]
             codes[i:i + 1] = replace_code
-        for i in enumerate(codes):
-            print(i)
+            arg_count = codes[j + 1].arg
+            replace_code_2 = [
+                Instr('CALL', arg_count + 1)
+            ]
+            codes[j + 1:j + 2] = replace_code_2
         func.__code__ = codes.to_code()
         return func
-
 
-    def hook_call(func_name: str, *args, **kwargs):
+    def hook_call(func_name, *args, **kwargs):
+        print(func_name)
+        print(args)
+        print(kwargs)
         print(f"Hooked {func_name} with {args} and {kwargs}")
 
+    def set_global(self, name, attr):
+        # This is for using it in every global scope
+        for frame in inspect.stack():
+            frame[0].f_globals[name] = attr
+
     def __init__(self, func):
-        globals()['__hook_func__'] = HookGlobals.hook_call # 不丢到 globals 里不知道怎么调用
-        self.replace_global_call(func)
+        self.set_global("__hook_func__", HookGlobals.hook_call)
+        if not hasattr(func, 'hooked__'):
+            self.replace_global_call(func)
+            func.__hooked__ = True
diff --git a/python/tester/src/direct_script.py b/python/tester/src/direct_script.py
diff --git a/src/exec/cli_exec.rs b/src/exec/cli_exec.rs
@@ -6,10 +6,7 @@ use std::{
 };
 
 use crate::{
-    consts::DURATION,
-    err,
-    term::tty::{DynTty, Tty, WrapperTty},
-    util::{anybase::AnyBase, util::rand_string},
+    consts::DURATION, err, log, term::tty::{DynTty, Tty, WrapperTty}, util::{anybase::AnyBase, util::rand_string}
 };
 
 use super::cli_api::{CliTestApi, ExecBase};
@@ -26,6 +23,7 @@ impl CliTester {
 
 impl CliTester {
     fn run_command(&mut self, command: &String) -> Result<(), Box<dyn Error>> {
+        log!("Write to shell: {}", command);
         let res = self.inner.write(command.as_bytes());
         if let Err(e) = res {
             return Err(e);

diff --git a/src/exec/runner.rs b/src/exec/runner.rs
@@ -1,35 +1,38 @@
 use super::cli_api::SudoCliTestApi;
 
-
 pub struct Cmd {
-    cmd: String
+    cmd: String,
 }
 
 pub struct TimeoutCmd {
     cmd: String,
-    timeout: u32
+    timeout: u32,
+}
+
+pub struct WaitCmd {
+    cmd: String,
+    wait: String,
+    timeout: u32,
 }
 
 pub enum CmdType {
-    Direct(Cmd), // writeln(cmd)
-    Run(TimeoutCmd), // script_run(cmd, timeout)
-    AssertRun(TimeoutCmd), // assert_script_run(cmd, timeout)
-    SudoRun(TimeoutCmd), // script_sudo(cmd, timeout)
+    Direct(Cmd),               // writeln(cmd)
+    Run(TimeoutCmd),           // script_run(cmd, timeout)
+    AssertRun(TimeoutCmd),     // assert_script_run(cmd, timeout)
+    SudoRun(TimeoutCmd),       // script_sudo(cmd, timeout)
     SudoAssertRun(TimeoutCmd), // assert_script_sudo(cmd, timeout)
-    Wait(TimeoutCmd), // wait_serial(cmd, timeout)
+    Wait(TimeoutCmd),          // wait_serial(cmd, timeout)
+    WaitRun(WaitCmd),          // writeln(cmd) wait_serial(wait, timeout)
 }
 
 pub struct CmdRunner {
     cmds: Vec<CmdType>,
-    inner: Box<dyn SudoCliTestApi>
+    inner: Box<dyn SudoCliTestApi>,
 }
 
 impl CmdRunner {
     pub fn build(inner: Box<dyn SudoCliTestApi>, cmds: Vec<CmdType>) -> Self {
-        Self {
-            cmds,
-            inner
-        }
+        Self { cmds, inner }
     }
 }
 
@@ -51,22 +54,28 @@ impl CmdRunner {
             match cmd {
                 CmdType::Direct(cmd) => {
                     self.inner.writeln(&cmd.cmd).unwrap();
-                },
+                }
                 CmdType::Run(cmd) => {
                     self.inner.script_run(&cmd.cmd, cmd.timeout).unwrap();
-                },
+                }
                 CmdType::AssertRun(cmd) => {
                     self.inner.assert_script_run(&cmd.cmd, cmd.timeout).unwrap();
-                },
+                }
                 CmdType::SudoRun(cmd) => {
                     self.inner.script_sudo(&cmd.cmd, cmd.timeout).unwrap();
-                },
+                }
                 CmdType::SudoAssertRun(cmd) => {
-                    self.inner.assert_script_sudo(&cmd.cmd, cmd.timeout).unwrap();
-                },
+                    self.inner
+                        .assert_script_sudo(&cmd.cmd, cmd.timeout)
+                        .unwrap();
+                }
                 CmdType::Wait(cmd) => {
                     self.inner.wait_serial(&cmd.cmd, cmd.timeout).unwrap();
                 }
+                CmdType::WaitRun(cmd) => {
+                    self.inner.writeln(&cmd.cmd).unwrap();
+                    self.inner.wait_serial(&cmd.wait, cmd.timeout).unwrap();
+                }
             }
         }
     }

diff --git a/src/pythonapi/pyhook.rs b/src/pythonapi/pyhook.rs
diff --git a/src/term/shell.rs b/src/term/shell.rs
@@ -141,6 +141,9 @@ impl Tty for Shell {
         let mut buff = self.buff.lock().unwrap();
         res.extend(buff.iter());
         buff.clear();
+        if !res.is_empty() {
+            log!("Shell read: {:?}", String::from_utf8_lossy(&res));
+        }
         return Ok(res);
     }
     fn read_line(&mut self) -> Result<Vec<u8>, Box<dyn Error>> {