diff --git a/Makefile b/Makefile index 7d33c06..01b719c 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,6 @@ compile: chmod +x ./tools/compile.py - ./tools/compile.py ./examples/execve2.py + ./tools/compile.py ./examples/execve3.py install: pip install -e . diff --git a/examples/execve3.py b/examples/execve3.py index 4e53d43..675d0e1 100644 --- a/examples/execve3.py +++ b/examples/execve3.py @@ -1,9 +1,10 @@ from pythonbpf import bpf, map, section, bpfglobal, compile -from pythonbpf.helpers import bpf_ktime_get_ns +from pythonbpf.helpers import ktime from pythonbpf.maps import HashMap from ctypes import c_void_p, c_int64, c_int32, c_uint64 + @bpf @map def last() -> HashMap: @@ -24,13 +25,16 @@ def hello_again(ctx: c_void_p) -> c_int64: print("exited") key = 0 tsp = last().lookup(key) - if tsp: - delta = (bpf_ktime_get_ns() - tsp.value) - if delta < 1000000000: - print("execve called within last second") - last().delete(key) - ts = bpf_ktime_get_ns() - last().update(key, ts) +# if tsp: +# delta = (bpf_ktime_get_ns() - tsp.value) +# if delta < 1000000000: +# print("execve called within last second") +# last().delete(key) + if True: + print("we prevailed") +# ts = ktime() + ktime() +# last().update(key, ts) return c_int64(0) @@ -39,4 +43,5 @@ def hello_again(ctx: c_void_p) -> c_int64: def LICENSE() -> str: return "GPL" + compile() diff --git a/pythonbpf/functions_pass.py b/pythonbpf/functions_pass.py index b8d4a6c..7265133 100644 --- a/pythonbpf/functions_pass.py +++ b/pythonbpf/functions_pass.py @@ -22,7 +22,7 @@ def get_probe_string(func_node): return "helper" -def handle_assign(module, builder, stmt, map_sym_tab, local_sym_tab): +def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab): """Handle assignment statements in the function body.""" if len(stmt.targets) != 1: print("Unsupported multiassignment") @@ -74,48 +74,113 @@ def handle_assign(module, builder, stmt, map_sym_tab, local_sym_tab): map_ptr = map_sym_tab[map_name] if method_name in helper_func_list: handle_helper_call( - rval, module, builder, None, local_sym_tab, map_sym_tab) + rval, module, builder, func, local_sym_tab, map_sym_tab) else: print("Unsupported assignment call structure") else: print("Unsupported assignment call function type") -def handle_if_statement(module, builder, stmt, map_sym_tab, local_sym_tab): - pass - - -def handle_expr(module, builder, expr, local_sym_tab, map_sym_tab): +def handle_expr(func, module, builder, expr, local_sym_tab, map_sym_tab): """Handle expression statements in the function body.""" call = expr.value + print(f"Handling expression: {ast.dump(call)}") if isinstance(call, ast.Call): if isinstance(call.func, ast.Name): # check for helpers first if call.func.id in helper_func_list: handle_helper_call( - call, module, builder, None, local_sym_tab, map_sym_tab) + call, module, builder, func, local_sym_tab, map_sym_tab) return - print("Unsupported expression statement") + elif isinstance(call, ast.Name): + if call.id in local_sym_tab: + var = local_sym_tab[call.id] + val = builder.load(var) + return val + else: + print(f"Undefined variable {call.id}") + return None + elif isinstance(call, ast.Constant): + if isinstance(call.value, int): + return ir.Constant(ir.IntType(64), call.value) + elif isinstance(call.value, bool): + return ir.Constant(ir.IntType(1), int(call.value)) + else: + print("Unsupported constant type") + return None + else: + print("Unsupported expression statement") -def handle_if(module, builder, stmt, map_sym_tab, local_sym_tab): +def handle_cond(func, module, builder, cond, local_sym_tab, map_sym_tab): + if isinstance(cond, ast.Constant): + if isinstance(cond.value, bool): + return ir.Constant(ir.IntType(1), int(cond.value)) + elif isinstance(cond.value, int): + return ir.Constant(ir.IntType(1), int(bool(cond.value))) + else: + print("Unsupported constant type in condition") + return None + elif isinstance(cond, ast.Name): + if cond.id in local_sym_tab: + var = local_sym_tab[cond.id] + val = builder.load(var) + return val + else: + print(f"Undefined variable {cond.id} in condition") + return None + else: + print("Unsupported condition expression") + return None + + +def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab): """Handle if statements in the function body.""" - func = builder.block.parent + print("Handling if statement") + start = builder.block.parent then_block = func.append_basic_block(name="if.then") merge_block = func.append_basic_block(name="if.end") - cond = stmt.test + cond = handle_cond(func, module, builder, stmt.test, + local_sym_tab, map_sym_tab) builder.cbranch(cond, then_block, merge_block) builder.position_at_end(then_block) for s in stmt.body: - pass + process_stmt(func, module, builder, s, + local_sym_tab, map_sym_tab, False) if not builder.block.is_terminated: builder.branch(merge_block) builder.position_at_end(merge_block) +def process_stmt(func, module, builder, stmt, local_sym_tab, map_sym_tab, did_return, ret_type=ir.IntType(64)): + print(f"Processing statement: {ast.dump(stmt)}") + if isinstance(stmt, ast.Expr): + handle_expr(func, module, builder, stmt, local_sym_tab, map_sym_tab) + elif isinstance(stmt, ast.Assign): + handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab) + elif isinstance(stmt, ast.If): + handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab) + elif isinstance(stmt, ast.Return): + if stmt.value is None: + builder.ret(ir.Constant(ir.IntType(32), 0)) + did_return = True + elif isinstance(stmt.value, ast.Call) and isinstance(stmt.value.func, ast.Name) and len(stmt.value.args) == 1 and isinstance(stmt.value.args[0], ast.Constant) and isinstance(stmt.value.args[0].value, int): + call_type = stmt.value.func.id + if ctypes_to_ir(call_type) != ret_type: + raise ValueError("Return type mismatch: expected" + f"{ctypes_to_ir(call_type)}, got {call_type}") + else: + builder.ret(ir.Constant( + ret_type, stmt.value.args[0].value)) + did_return = True + else: + print("Unsupported return value") + return did_return + + def process_func_body(module, builder, func_node, func, ret_type, map_sym_tab): """Process the body of a bpf function""" # TODO: A lot. We just have print -> bpf_trace_printk for now @@ -124,27 +189,9 @@ def process_func_body(module, builder, func_node, func, ret_type, map_sym_tab): local_sym_tab = {} for stmt in func_node.body: - if isinstance(stmt, ast.Expr): - handle_expr(module, builder, stmt, local_sym_tab, map_sym_tab) - elif isinstance(stmt, ast.Assign): - handle_assign(module, builder, stmt, map_sym_tab, local_sym_tab) - elif isinstance(stmt, ast.If): - handle_if(module, builder, stmt, map_sym_tab, local_sym_tab) - elif isinstance(stmt, ast.Return): - if stmt.value is None: - builder.ret(ir.Constant(ir.IntType(32), 0)) - did_return = True - elif isinstance(stmt.value, ast.Call) and isinstance(stmt.value.func, ast.Name) and len(stmt.value.args) == 1 and isinstance(stmt.value.args[0], ast.Constant) and isinstance(stmt.value.args[0].value, int): - call_type = stmt.value.func.id - if ctypes_to_ir(call_type) != ret_type: - raise ValueError("Return type mismatch: expected" - f"{ctypes_to_ir(call_type)}, got {call_type}") - else: - builder.ret(ir.Constant( - ret_type, stmt.value.args[0].value)) - did_return = True - else: - print("Unsupported return value") + did_return = process_stmt(func, module, builder, stmt, local_sym_tab, + map_sym_tab, did_return, ret_type) + if not did_return: builder.ret(ir.Constant(ir.IntType(32), 0))