Skip to content

Commit aeb9a45

Browse files
r41k0ur41k0u
committed
Add condition eval and basic if example - workin
1 parent 357ad7c commit aeb9a45

3 files changed

Lines changed: 95 additions & 43 deletions

File tree

‎Makefile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
compile:
22
chmod +x ./tools/compile.py
3-
./tools/compile.py ./examples/execve2.py
3+
./tools/compile.py ./examples/execve3.py
44

55
install:
66
pip install -e .

‎examples/execve3.py‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
from pythonbpf import bpf, map, section, bpfglobal, compile
2-
from pythonbpf.helpers import bpf_ktime_get_ns
2+
from pythonbpf.helpers import ktime
33
from pythonbpf.maps import HashMap
44

55
from ctypes import c_void_p, c_int64, c_int32, c_uint64
66

7+
78
@bpf
89
@map
910
def last() -> HashMap:
@@ -24,13 +25,16 @@ def hello_again(ctx: c_void_p) -> c_int64:
2425
print("exited")
2526
key = 0
2627
tsp = last().lookup(key)
27-
if tsp:
28-
delta = (bpf_ktime_get_ns() - tsp.value)
29-
if delta < 1000000000:
30-
print("execve called within last second")
31-
last().delete(key)
32-
ts = bpf_ktime_get_ns()
33-
last().update(key, ts)
28+
# if tsp:
29+
# delta = (bpf_ktime_get_ns() - tsp.value)
30+
# if delta < 1000000000:
31+
# print("execve called within last second")
32+
# last().delete(key)
33+
if True:
34+
print("we prevailed")
35+
# ts = ktime()
36+
ktime()
37+
# last().update(key, ts)
3438
return c_int64(0)
3539

3640

@@ -39,4 +43,5 @@ def hello_again(ctx: c_void_p) -> c_int64:
3943
def LICENSE() -> str:
4044
return "GPL"
4145

46+
4247
compile()

‎pythonbpf/functions_pass.py‎

Lines changed: 81 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ def get_probe_string(func_node):
2222
return "helper"
2323

2424

25-
def handle_assign(module, builder, stmt, map_sym_tab, local_sym_tab):
25+
def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab):
2626
"""Handle assignment statements in the function body."""
2727
if len(stmt.targets) != 1:
2828
print("Unsupported multiassignment")
@@ -74,48 +74,113 @@ def handle_assign(module, builder, stmt, map_sym_tab, local_sym_tab):
7474
map_ptr = map_sym_tab[map_name]
7575
if method_name in helper_func_list:
7676
handle_helper_call(
77-
rval, module, builder, None, local_sym_tab, map_sym_tab)
77+
rval, module, builder, func, local_sym_tab, map_sym_tab)
7878
else:
7979
print("Unsupported assignment call structure")
8080
else:
8181
print("Unsupported assignment call function type")
8282

8383

84-
def handle_if_statement(module, builder, stmt, map_sym_tab, local_sym_tab):
85-
pass
86-
87-
88-
def handle_expr(module, builder, expr, local_sym_tab, map_sym_tab):
84+
def handle_expr(func, module, builder, expr, local_sym_tab, map_sym_tab):
8985
"""Handle expression statements in the function body."""
9086
call = expr.value
87+
print(f"Handling expression: {ast.dump(call)}")
9188
if isinstance(call, ast.Call):
9289
if isinstance(call.func, ast.Name):
9390
# check for helpers first
9491
if call.func.id in helper_func_list:
9592
handle_helper_call(
96-
call, module, builder, None, local_sym_tab, map_sym_tab)
93+
call, module, builder, func, local_sym_tab, map_sym_tab)
9794
return
98-
print("Unsupported expression statement")
95+
elif isinstance(call, ast.Name):
96+
if call.id in local_sym_tab:
97+
var = local_sym_tab[call.id]
98+
val = builder.load(var)
99+
return val
100+
else:
101+
print(f"Undefined variable {call.id}")
102+
return None
103+
elif isinstance(call, ast.Constant):
104+
if isinstance(call.value, int):
105+
return ir.Constant(ir.IntType(64), call.value)
106+
elif isinstance(call.value, bool):
107+
return ir.Constant(ir.IntType(1), int(call.value))
108+
else:
109+
print("Unsupported constant type")
110+
return None
111+
else:
112+
print("Unsupported expression statement")
113+
114+
115+
def handle_cond(func, module, builder, cond, local_sym_tab, map_sym_tab):
116+
if isinstance(cond, ast.Constant):
117+
if isinstance(cond.value, bool):
118+
return ir.Constant(ir.IntType(1), int(cond.value))
119+
elif isinstance(cond.value, int):
120+
return ir.Constant(ir.IntType(1), int(bool(cond.value)))
121+
else:
122+
print("Unsupported constant type in condition")
123+
return None
124+
elif isinstance(cond, ast.Name):
125+
if cond.id in local_sym_tab:
126+
var = local_sym_tab[cond.id]
127+
val = builder.load(var)
128+
return val
129+
else:
130+
print(f"Undefined variable {cond.id} in condition")
131+
return None
132+
else:
133+
print("Unsupported condition expression")
134+
return None
99135

100136

101-
def handle_if(module, builder, stmt, map_sym_tab, local_sym_tab):
137+
def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab):
102138
"""Handle if statements in the function body."""
103-
func = builder.block.parent
139+
print("Handling if statement")
140+
start = builder.block.parent
104141
then_block = func.append_basic_block(name="if.then")
105142
merge_block = func.append_basic_block(name="if.end")
106143

107-
cond = stmt.test
144+
cond = handle_cond(func, module, builder, stmt.test,
145+
local_sym_tab, map_sym_tab)
108146

109147
builder.cbranch(cond, then_block, merge_block)
110148
builder.position_at_end(then_block)
111149
for s in stmt.body:
112-
pass
150+
process_stmt(func, module, builder, s,
151+
local_sym_tab, map_sym_tab, False)
113152
if not builder.block.is_terminated:
114153
builder.branch(merge_block)
115154

116155
builder.position_at_end(merge_block)
117156

118157

158+
def process_stmt(func, module, builder, stmt, local_sym_tab, map_sym_tab, did_return, ret_type=ir.IntType(64)):
159+
print(f"Processing statement: {ast.dump(stmt)}")
160+
if isinstance(stmt, ast.Expr):
161+
handle_expr(func, module, builder, stmt, local_sym_tab, map_sym_tab)
162+
elif isinstance(stmt, ast.Assign):
163+
handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab)
164+
elif isinstance(stmt, ast.If):
165+
handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab)
166+
elif isinstance(stmt, ast.Return):
167+
if stmt.value is None:
168+
builder.ret(ir.Constant(ir.IntType(32), 0))
169+
did_return = True
170+
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):
171+
call_type = stmt.value.func.id
172+
if ctypes_to_ir(call_type) != ret_type:
173+
raise ValueError("Return type mismatch: expected"
174+
f"{ctypes_to_ir(call_type)}, got {call_type}")
175+
else:
176+
builder.ret(ir.Constant(
177+
ret_type, stmt.value.args[0].value))
178+
did_return = True
179+
else:
180+
print("Unsupported return value")
181+
return did_return
182+
183+
119184
def process_func_body(module, builder, func_node, func, ret_type, map_sym_tab):
120185
"""Process the body of a bpf function"""
121186
# 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):
124189
local_sym_tab = {}
125190

126191
for stmt in func_node.body:
127-
if isinstance(stmt, ast.Expr):
128-
handle_expr(module, builder, stmt, local_sym_tab, map_sym_tab)
129-
elif isinstance(stmt, ast.Assign):
130-
handle_assign(module, builder, stmt, map_sym_tab, local_sym_tab)
131-
elif isinstance(stmt, ast.If):
132-
handle_if(module, builder, stmt, map_sym_tab, local_sym_tab)
133-
elif isinstance(stmt, ast.Return):
134-
if stmt.value is None:
135-
builder.ret(ir.Constant(ir.IntType(32), 0))
136-
did_return = True
137-
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):
138-
call_type = stmt.value.func.id
139-
if ctypes_to_ir(call_type) != ret_type:
140-
raise ValueError("Return type mismatch: expected"
141-
f"{ctypes_to_ir(call_type)}, got {call_type}")
142-
else:
143-
builder.ret(ir.Constant(
144-
ret_type, stmt.value.args[0].value))
145-
did_return = True
146-
else:
147-
print("Unsupported return value")
192+
did_return = process_stmt(func, module, builder, stmt, local_sym_tab,
193+
map_sym_tab, did_return, ret_type)
194+
148195
if not did_return:
149196
builder.ret(ir.Constant(ir.IntType(32), 0))
150197

0 commit comments

Comments
 (0)