bump version

Signed-off-by: varun-r-mallya <varunrmallya@gmail.com>

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 .

README.md

@@ -1,15 +1,70 @@
 # Python-BPF
-This is an LLVM IR generator for eBPF programs in Python. We use `llvmlite` to generate LLVM IR code from pure Python code. This is then compiled to LLVM object files, which can be loaded into the kernel for execution.
+<p align="center">
+<a href="https://www.python.org/downloads/release/python-3080/"><img src="https://img.shields.io/badge/python-3.8-blue.svg"></a>
+<a href="https://pypi.org/project/pythonbpf"><img src="https://badge.fury.io/py/pythonbpf.svg"></a>
+</p>
+
+This is an LLVM IR generator for eBPF programs in Python. We use llvmlite to generate LLVM IR from pure Python. This is then compiled to LLVM object files, which can be loaded into the kernel for execution. We do not rely on BCC to do our compilation.
+
+# DO NOT USE IN PRODUCTION. IN DEVELOPMENT.
+
+##  Installation
+- Have `clang` installed.
+- `pip install pythonbpf`
+
+## Usage
+```python
+# pythonbpf_example.py
+from pythonbpf import bpf, map, bpfglobal, section, compile
+from pythonbpf.helpers import bpf_ktime_get_ns
+from pythonbpf.maps import HashMap
+
+from ctypes import c_void_p, c_int64, c_int32, c_uint64
+
+@bpf
+@map
+def last() -> HashMap:
+    return HashMap(key_type=c_uint64, value_type=c_uint64, max_entries=1)
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_execve")
+def hello(ctx: c_void_p) -> c_int32:
+    print("entered")
+    return c_int32(0)
+
+@bpf
+@section("tracepoint/syscalls/sys_exit_execve")
+def hello_again(ctx: c_void_p) -> c_int64:
+    print("exited")
+    key = 0
+    tsp = last().lookup(key)
+    print(tsp)
+    ts = bpf_ktime_get_ns()
+    return c_int64(0)
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+def some_normal_function():
+    print("normal function")
+
+# compiles and dumps object file in the same directory
+compile()
+```
+- Run `python pythonbpf_example.py` to get the compiled object file that can be then loaded into the kernel.
 
 ## Development
-Step 0. Make a virtual environment and activate it using `python3 -m venv .venv && source .venv/bin/activate`.  
-Step 1. Run `make install` to install the required dependencies.  
-Step 2. Run `make` to see the compilation output of the example.  
-Step 3. Run `check.sh` to check if generated object file passes through the verifier inside the examples directory.  
-Step 4. Run `make` in the `examples/c-form` directory to modify the example C BPF program to check the actual LLVM IR generated by clang.  
+- Make a virtual environment and activate it using `python3 -m venv .venv && source .venv/bin/activate`.  
+- Run `make install` to install the required dependencies.  
+- Run `make` to see the compilation output of the example.  
+- Run `check.sh` to check if generated object file passes through the verifier inside the examples directory.  
+- Run `make` in the `examples/c-form` directory to modify the example C BPF program to check the actual LLVM IR generated by clang.
 
 ### Development Notes
-Run ` ./check.sh run execve2.o;` in examples folder
+- Run ` ./check.sh check execve2.o;` in examples folder to check if the object code passes the verifier.
+- Run ` ./check.sh run execve2.o;` in examples folder to run the object code using `bpftool`.
 
 ## Authors
 - [@r41k0u](https://github.com/r41k0u)

examples/c-form/ex3-2.bpf.c

@@ -8,7 +8,7 @@ struct {
     __uint(type, BPF_MAP_TYPE_HASH);
     __type(key, u64);
     __type(value, u64);
-    __uint(max_entries, 1);
+    __uint(max_entries, 4);
 } last SEC(".maps");
 
 // Handler for syscall entry

examples/c-form/ex4.bpf.c

@@ -33,7 +33,7 @@ long hello_again(void *ctx) {
         u64 delta = bpf_ktime_get_ns() - *tsp;
         if (delta < 1000000000) {
             // output if time is less than 1 second
-            bpf_trace_printk("%d\\n", delta / 1000000);
+            bpf_printk("execve called within last second");
         }
         bpf_map_delete_elem(&last, &key);
     }

examples/execve2.py

@@ -1,6 +1,6 @@
 from pythonbpf.decorators import bpf, map, section, bpfglobal
 from ctypes import c_void_p, c_int64, c_int32, c_uint64
-from pythonbpf.helpers import bpf_ktime_get_ns
+from pythonbpf.helpers import ktime
 from pythonbpf.maps import HashMap
 
 
@@ -9,6 +9,7 @@ from pythonbpf.maps import HashMap
 def last() -> HashMap:
     return HashMap(key_type=c_uint64, value_type=c_uint64, max_entries=1)
 
+
 @bpf
 @section("tracepoint/syscalls/sys_enter_execve")
 def hello(ctx: c_void_p) -> c_int32:
@@ -24,9 +25,10 @@ def hello_again(ctx: c_void_p) -> c_int64:
     key = 0
     tsp = last().lookup(key)
     print(tsp)
-    ts = bpf_ktime_get_ns()
+    ts = ktime()
     return c_int64(0)
 
+
 @bpf
 @bpfglobal
 def LICENSE() -> str:

examples/execve3.py

@@ -1,13 +1,14 @@
-from pythonbpf.decorators import bpf, map, section, bpfglobal
-from ctypes import c_void_p, c_int64, c_int32, c_uint64
-from pythonbpf.helpers import bpf_ktime_get_ns
+from pythonbpf import bpf, map, section, bpfglobal, compile
+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:
-    return HashMap(key_type=c_uint64, value_type=c_uint64, max_entries=1)
+    return HashMap(key_type=c_uint64, value_type=c_uint64, max_entries=3)
 
 
 @bpf
@@ -24,17 +25,36 @@ 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()
+#    if tsp:
+#        delta = (bpf_ktime_get_ns() - tsp.value)
+#        if delta < 1000000000:
+#            print("execve called within last second")
+#        last().delete(key)
+    x = 1
+    y = False
+    if x > 0:
+        if x < 2:
+            print(f"we prevailed {x}")
+        else:
+            print(f"we did not prevail {x}")
+    ts = ktime()
     last().update(key, ts)
-    return c_int64(0)
 
+    st = "st"
+    last().update(key, ts)
+
+    keena = 2 + 1
+    # below breaks
+    # keela = keena + 1
+    # TODO: binops evaluate but into a random register and dont get assigned.
+    keema = 8 * 9
+    keesa = 10 - 11
+    keeda = 10 / 5
+    return c_int64(0)
 
 @bpf
 @bpfglobal
 def LICENSE() -> str:
     return "GPL"
+
+compile()

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pythonbpf"
-version = "0.1.0"
+version = "0.1.1"
 description = "Reduced Python frontend for eBPF"
 authors = [
   { name = "r41k0u", email="pragyanshchaturvedi18@gmail.com" },

pythonbpf/__init__.py

@@ -0,0 +1,2 @@
+from .decorators import bpf, map, section, bpfglobal
+from .codegen import compile_to_ir, compile

pythonbpf/binary_ops.py

@@ -0,0 +1,35 @@
+import ast
+from llvmlite import ir
+
+def handle_binary_op(rval, module, builder, func, local_sym_tab, map_sym_tab):
+    left = rval.left
+    right = rval.right
+    op = rval.op
+
+    if isinstance(left, ast.Name):
+        left = local_sym_tab[left.id]
+    elif isinstance(left, ast.Constant):
+        left = ir.Constant(ir.IntType(64), left.value)
+    else:
+        print("Unsupported left operand type")
+
+    if isinstance(right, ast.Name):
+        right = local_sym_tab[right.id]
+    elif isinstance(right, ast.Constant):
+        right = ir.Constant(ir.IntType(64), right.value)
+    else:
+        SyntaxError("Unsupported right operand type")
+
+    if isinstance(op, ast.Add):
+        result = builder.add(left, right)
+    elif isinstance(op, ast.Sub):
+        result = builder.sub(left, right)
+    elif isinstance(op, ast.Mult):
+        result = builder.mul(left, right)
+    elif isinstance(op, ast.Div):
+        result = builder.sdiv(left, right)
+    else:
+        result = "fuck type errors"
+        SyntaxError("Unsupported binary operation")
+
+    return result

pythonbpf/bpf_helper_handler.py

@@ -1,20 +1,52 @@
 import ast
 from llvmlite import ir
+from .expr_pass import eval_expr
 
 
-def bpf_ktime_get_ns_emitter(call, module, builder, func):
-    pass
+def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab=None):
+    """
+    Emit LLVM IR for bpf_ktime_get_ns helper function call.
+    """
+    # func is an arg to just have a uniform signature with other emitters
+    helper_id = ir.Constant(ir.IntType(64), 5)
+    fn_type = ir.FunctionType(ir.IntType(64), [], var_arg=False)
+    fn_ptr_type = ir.PointerType(fn_type)
+    fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
+    result = builder.call(fn_ptr, [], tail=False)
+    return result
 
 
-def bpf_map_lookup_elem_emitter(map_ptr, key_ptr, module, builder):
+def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, local_sym_tab=None):
     """
     Emit LLVM IR for bpf_map_lookup_elem helper function call.
     """
-    # Cast pointers to void*
+    if call.args and len(call.args) != 1:
+        raise ValueError("Map lookup expects exactly one argument, got "
+                         f"{len(call.args)}")
+    key_arg = call.args[0]
+    if isinstance(key_arg, ast.Name):
+        key_name = key_arg.id
+        if local_sym_tab and key_name in local_sym_tab:
+            key_ptr = local_sym_tab[key_name]
+        else:
+            raise ValueError(
+                f"Key variable {key_name} not found in local symbol table.")
+    elif isinstance(key_arg, ast.Constant) and isinstance(key_arg.value, int):
+        # handle constant integer keys
+        key_val = key_arg.value
+        key_type = ir.IntType(64)
+        key_ptr = builder.alloca(key_type)
+        key_ptr.align = key_type // 8
+        builder.store(ir.Constant(key_type, key_val), key_ptr)
+    else:
+        raise NotImplementedError(
+            "Only simple variable names are supported as keys in map lookup.")
+
+    if key_ptr is None:
+        raise ValueError("Key pointer is None.")
+
     map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
 
-    # Define function type for bpf_map_lookup_elem
-    # The function takes two void* arguments and returns void*
     fn_type = ir.FunctionType(
         ir.PointerType(),  # Return type: void*
         [ir.PointerType(), ir.PointerType()],  # Args: (void*, void*)
@@ -26,16 +58,92 @@ def bpf_map_lookup_elem_emitter(map_ptr, key_ptr, module, builder):
     fn_addr = ir.Constant(ir.IntType(64), 1)
     fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
 
-    # Call the helper function
     result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)
 
     return result
 
 
-def bpf_printk_emitter(call, module, builder, func):
+def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None):
     if not hasattr(func, "_fmt_counter"):
         func._fmt_counter = 0
 
+    if not call.args:
+        raise ValueError("print expects at least one argument")
+
+    if isinstance(call.args[0], ast.JoinedStr):
+        fmt_parts = []
+        exprs = []
+
+        for value in call.args[0].values:
+            if isinstance(value, ast.Constant):
+                if isinstance(value.value, str):
+                    fmt_parts.append(value.value)
+                elif isinstance(value.value, int):
+                    fmt_parts.append("%lld")
+                    exprs.append(ir.Constant(ir.IntType(64), value.value))
+                else:
+                    raise NotImplementedError(
+                        "Only string and integer constants are supported in f-string.")
+            elif isinstance(value, ast.FormattedValue):
+                # Assume int for now
+                fmt_parts.append("%d")
+                if isinstance(value.value, ast.Name):
+                    exprs.append(value.value)
+                else:
+                    raise NotImplementedError(
+                        "Only simple variable names are supported in formatted values.")
+            else:
+                raise NotImplementedError(
+                    "Unsupported value type in f-string.")
+
+        fmt_str = "".join(fmt_parts) + "\n" + "\0"
+        fmt_name = f"{func.name}____fmt{func._fmt_counter}"
+        func._fmt_counter += 1
+
+        fmt_gvar = ir.GlobalVariable(
+            module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=fmt_name)
+        fmt_gvar.global_constant = True
+        fmt_gvar.initializer = ir.Constant(
+            ir.ArrayType(ir.IntType(8), len(fmt_str)),
+            bytearray(fmt_str.encode("utf8"))
+        )
+        fmt_gvar.linkage = "internal"
+        fmt_gvar.align = 1
+
+        fmt_ptr = builder.bitcast(fmt_gvar, ir.PointerType())
+
+        args = [fmt_ptr, ir.Constant(ir.IntType(32), len(fmt_str))]
+
+        # Only 3 args supported in bpf_printk
+        if len(exprs) > 3:
+            print(
+                "Warning: bpf_printk supports up to 3 arguments, extra arguments will be ignored.")
+
+        for expr in exprs[:3]:
+            val = eval_expr(func, module, builder, expr, local_sym_tab, None)
+            if val:
+                if isinstance(val.type, ir.PointerType):
+                    val = builder.ptrtoint(val, ir.IntType(64))
+                elif isinstance(val.type, ir.IntType):
+                    if val.type.width < 64:
+                        val = builder.sext(val, ir.IntType(64))
+                else:
+                    print(
+                        "Warning: Only integer and pointer types are supported in bpf_printk arguments. Others will be converted to 0.")
+                    val = ir.Constant(ir.IntType(64), 0)
+                args.append(val)
+            else:
+                print(
+                    "Warning: Failed to evaluate expression for bpf_printk argument. It will be converted to 0.")
+                args.append(ir.Constant(ir.IntType(64), 0))
+
+        fn_type = ir.FunctionType(ir.IntType(
+            64), [ir.PointerType(), ir.IntType(32)], var_arg=True)
+        fn_ptr_type = ir.PointerType(fn_type)
+        fn_addr = ir.Constant(ir.IntType(64), 6)
+        fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
+        return builder.call(fn_ptr, args, tail=True)
+
     for arg in call.args:
         if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
             fmt_str = arg.value + "\n" + "\0"
@@ -62,3 +170,137 @@ def bpf_printk_emitter(call, module, builder, func):
 
             builder.call(fn_ptr, [fmt_ptr, ir.Constant(
                 ir.IntType(32), len(fmt_str))], tail=True)
+
+
+def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=None):
+    """
+    Emit LLVM IR for bpf_map_update_elem helper function call.
+    Expected call signature: map.update(key, value, flags=0)
+    """
+    if not call.args or len(call.args) < 2 or len(call.args) > 3:
+        raise ValueError("Map update expects 2 or 3 arguments (key, value, flags), got "
+                         f"{len(call.args)}")
+
+    key_arg = call.args[0]
+    value_arg = call.args[1]
+    flags_arg = call.args[2] if len(call.args) > 2 else None
+
+    # Handle key
+    if isinstance(key_arg, ast.Name):
+        key_name = key_arg.id
+        if local_sym_tab and key_name in local_sym_tab:
+            key_ptr = local_sym_tab[key_name]
+        else:
+            raise ValueError(
+                f"Key variable {key_name} not found in local symbol table.")
+    elif isinstance(key_arg, ast.Constant) and isinstance(key_arg.value, int):
+        # Handle constant integer keys
+        key_val = key_arg.value
+        key_type = ir.IntType(64)
+        key_ptr = builder.alloca(key_type)
+        key_ptr.align = key_type.width // 8
+        builder.store(ir.Constant(key_type, key_val), key_ptr)
+    else:
+        raise NotImplementedError(
+            "Only simple variable names and integer constants are supported as keys in map update.")
+
+    # Handle value
+    if isinstance(value_arg, ast.Name):
+        value_name = value_arg.id
+        if local_sym_tab and value_name in local_sym_tab:
+            value_ptr = local_sym_tab[value_name]
+        else:
+            raise ValueError(
+                f"Value variable {value_name} not found in local symbol table.")
+    elif isinstance(value_arg, ast.Constant) and isinstance(value_arg.value, int):
+        # Handle constant integers
+        value_val = value_arg.value
+        value_type = ir.IntType(64)
+        value_ptr = builder.alloca(value_type)
+        value_ptr.align = value_type.width // 8
+        builder.store(ir.Constant(value_type, value_val), value_ptr)
+    else:
+        raise NotImplementedError(
+            "Only simple variable names and integer constants are supported as values in map update.")
+
+    # Handle flags argument (defaults to 0)
+    if flags_arg is not None:
+        if isinstance(flags_arg, ast.Constant) and isinstance(flags_arg.value, int):
+            flags_val = flags_arg.value
+        elif isinstance(flags_arg, ast.Name):
+            flags_name = flags_arg.id
+            if local_sym_tab and flags_name in local_sym_tab:
+                # Assume it's a stored integer value, load it
+                flags_ptr = local_sym_tab[flags_name]
+                flags_val = builder.load(flags_ptr)
+            else:
+                raise ValueError(
+                    f"Flags variable {flags_name} not found in local symbol table.")
+        else:
+            raise NotImplementedError(
+                "Only integer constants and simple variable names are supported as flags in map update.")
+    else:
+        flags_val = 0
+
+    if key_ptr is None or value_ptr is None:
+        raise ValueError("Key pointer or value pointer is None.")
+
+    map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
+    fn_type = ir.FunctionType(
+        ir.IntType(64),
+        [ir.PointerType(), ir.PointerType(), ir.PointerType(), ir.IntType(64)],
+        var_arg=False
+    )
+    fn_ptr_type = ir.PointerType(fn_type)
+
+    # helper id
+    fn_addr = ir.Constant(ir.IntType(64), 2)
+    fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
+
+    if isinstance(flags_val, int):
+        flags_const = ir.Constant(ir.IntType(64), flags_val)
+    else:
+        flags_const = flags_val
+
+    result = builder.call(
+        fn_ptr, [map_void_ptr, key_ptr, value_ptr, flags_const], tail=False)
+
+    return result
+
+
+helper_func_list = {
+    "lookup": bpf_map_lookup_elem_emitter,
+    "print": bpf_printk_emitter,
+    "ktime": bpf_ktime_get_ns_emitter,
+    "update": bpf_map_update_elem_emitter,
+}
+
+
+def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_tab=None):
+    if isinstance(call.func, ast.Name):
+        func_name = call.func.id
+        if func_name in helper_func_list:
+            # it is not a map method call
+            return helper_func_list[func_name](call, None, module, builder, func, local_sym_tab)
+        else:
+            raise NotImplementedError(
+                f"Function {func_name} is not implemented as a helper function.")
+    elif isinstance(call.func, ast.Attribute):
+        # likely a map method call
+        if isinstance(call.func.value, ast.Call) and isinstance(call.func.value.func, ast.Name):
+            map_name = call.func.value.func.id
+            method_name = call.func.attr
+            if map_sym_tab and map_name in map_sym_tab:
+                map_ptr = map_sym_tab[map_name]
+                if method_name in helper_func_list:
+                    return helper_func_list[method_name](
+                        call, map_ptr, module, builder, local_sym_tab)
+                else:
+                    raise NotImplementedError(
+                        f"Map method {method_name} is not implemented as a helper function.")
+            else:
+                raise ValueError(
+                    f"Map variable {map_name} not found in symbol tables.")
+        else:
+            raise NotImplementedError(
+                "Attribute not supported for map method calls.")

pythonbpf/codegen.py

@@ -5,7 +5,9 @@ from .functions_pass import func_proc
 from .maps_pass import maps_proc
 from .globals_pass import globals_processing
 import os
-
+import subprocess
+import inspect
+from pathlib import Path
 
 def find_bpf_chunks(tree):
     """Find all functions decorated with @bpf in the AST."""
@@ -92,3 +94,20 @@ def compile_to_ir(filename: str, output: str):
         f.write("\n")
 
     return output
+
+def compile():
+    # Look one level up the stack to the caller of this function
+    caller_frame = inspect.stack()[1]
+    caller_file = Path(caller_frame.filename).resolve()
+
+    ll_file = Path("/tmp") / caller_file.with_suffix(".ll").name
+    o_file = caller_file.with_suffix(".o")
+
+    compile_to_ir(str(caller_file), str(ll_file))
+
+    subprocess.run([
+        "llc", "-march=bpf", "-filetype=obj", "-O2",
+        str(ll_file), "-o", str(o_file)
+    ], check=True)
+
+    print(f"Object written to {o_file}")

pythonbpf/expr_pass.py

@@ -0,0 +1,49 @@
+import ast
+from llvmlite import ir
+
+
+def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab):
+    print(f"Evaluating expression: {expr}")
+    if isinstance(expr, ast.Name):
+        if expr.id in local_sym_tab:
+            var = local_sym_tab[expr.id]
+            val = builder.load(var)
+            return val
+        else:
+            print(f"Undefined variable {expr.id}")
+            return None
+    elif isinstance(expr, ast.Constant):
+        if isinstance(expr.value, int):
+            return ir.Constant(ir.IntType(64), expr.value)
+        elif isinstance(expr.value, bool):
+            return ir.Constant(ir.IntType(1), int(expr.value))
+        else:
+            print("Unsupported constant type")
+            return None
+    elif isinstance(expr, ast.Call):
+        # delayed import to avoid circular dependency
+        from .bpf_helper_handler import helper_func_list, handle_helper_call
+
+        if isinstance(expr.func, ast.Name):
+            # check for helpers first
+            if expr.func.id in helper_func_list:
+                return handle_helper_call(
+                    expr, module, builder, func, local_sym_tab, map_sym_tab)
+        elif isinstance(expr.func, ast.Attribute):
+            if isinstance(expr.func.value, ast.Call) and isinstance(expr.func.value.func, ast.Name):
+                method_name = expr.func.attr
+                if method_name in helper_func_list:
+                    return handle_helper_call(
+                        expr, module, builder, func, local_sym_tab, map_sym_tab)
+    print("Unsupported expression evaluation")
+    return None
+
+
+def handle_expr(func, module, builder, expr, local_sym_tab, map_sym_tab):
+    """Handle expression statements in the function body."""
+    print(f"Handling expression: {ast.dump(expr)}")
+    call = expr.value
+    if isinstance(call, ast.Call):
+        eval_expr(func, module, builder, call, local_sym_tab, map_sym_tab)
+    else:
+        print("Unsupported expression type")

pythonbpf/functions_pass.py

@@ -1,8 +1,11 @@
 from llvmlite import ir
 import ast
 
-from .bpf_helper_handler import bpf_printk_emitter, bpf_ktime_get_ns_emitter, bpf_map_lookup_elem_emitter
+
+from .bpf_helper_handler import helper_func_list, handle_helper_call
 from .type_deducer import ctypes_to_ir
+from .binary_ops import handle_binary_op
+from .expr_pass import eval_expr, handle_expr
 
 
 def get_probe_string(func_node):
@@ -22,7 +25,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")
@@ -37,78 +40,191 @@ def handle_assign(module, builder, stmt, map_sym_tab, local_sym_tab):
     var_name = target.id
     rval = stmt.value
     if isinstance(rval, ast.Constant):
-        if isinstance(rval.value, int):
-            # Assume c_int64 for now
-            # TODO: make symtab for this
-            var = builder.alloca(ir.IntType(64), name=var_name)
-            var.align = 8
-            builder.store(ir.Constant(ir.IntType(64), rval.value), var)
-            local_sym_tab[var_name] = var
+        if isinstance(rval.value, bool):
+            if rval.value:
+                builder.store(ir.Constant(ir.IntType(1), 1),
+                              local_sym_tab[var_name])
+            else:
+                builder.store(ir.Constant(ir.IntType(1), 0),
+                              local_sym_tab[var_name])
             print(f"Assigned constant {rval.value} to {var_name}")
+        elif isinstance(rval.value, int):
+            # Assume c_int64 for now
+            # var = builder.alloca(ir.IntType(64), name=var_name)
+            # var.align = 8
+            builder.store(ir.Constant(ir.IntType(64), rval.value),
+                          local_sym_tab[var_name])
+            # local_sym_tab[var_name] = var
+            print(f"Assigned constant {rval.value} to {var_name}")
+        else:
+            print("Unsupported constant type")
     elif isinstance(rval, ast.Call):
         if isinstance(rval.func, ast.Name):
             call_type = rval.func.id
             print(f"Assignment call type: {call_type}")
             if call_type in num_types and len(rval.args) == 1 and isinstance(rval.args[0], ast.Constant) and isinstance(rval.args[0].value, int):
                 ir_type = ctypes_to_ir(call_type)
-                var = builder.alloca(ir_type, name=var_name)
-                var.align = ir_type.width // 8
-                builder.store(ir.Constant(ir_type, rval.args[0].value), var)
+                # var = builder.alloca(ir_type, name=var_name)
+                # var.align = ir_type.width // 8
+                builder.store(ir.Constant(
+                    ir_type, rval.args[0].value), local_sym_tab[var_name])
                 print(f"Assigned {call_type} constant "
                       f"{rval.args[0].value} to {var_name}")
-                local_sym_tab[var_name] = var
+                # local_sym_tab[var_name] = var
+            elif call_type in helper_func_list:
+                # var = builder.alloca(ir.IntType(64), name=var_name)
+                # var.align = 8
+                val = handle_helper_call(
+                    rval, module, builder, None, local_sym_tab, map_sym_tab)
+                builder.store(val, local_sym_tab[var_name])
+                # local_sym_tab[var_name] = var
+                print(f"Assigned constant {rval.func.id} to {var_name}")
             else:
                 print(f"Unsupported assignment call type: {call_type}")
         elif isinstance(rval.func, ast.Attribute):
-            if isinstance(rval.func.attr, str) and rval.func.attr == "lookup":
-                # Get map name and check symtab
-                # maps are called as funcs
-                if isinstance(rval.func.value, ast.Call) and isinstance(rval.func.value.func, ast.Name):
-                    map_name = rval.func.value.func.id
-                    if map_name in map_sym_tab:
-                        map_global = map_sym_tab[map_name]
-                        print(f"Found map {map_name} in symtab for lookup")
-                        if len(rval.args) != 1:
-                            print("Unsupported lookup with != 1 arg")
-                            return
-                        key_arg = rval.args[0]
-                        print(f"Lookup key arg type: {type(key_arg)}")
-                        # TODO: implement a parse_arg ffs as this can be a fucking expr
-                        if isinstance(key_arg, ast.Constant) and isinstance(key_arg.value, int):
-                            key_val = key_arg.value
-                            key_type = ir.IntType(64)
-                            print(f"Key type: {key_type}")
-                            print(f"Key val: {key_val}")
-                            key_var = builder.alloca(key_type)
-                            key_var.align = key_type // 8
-                            builder.store(ir.Constant(
-                                key_type, key_val), key_var)
-                        elif isinstance(key_arg, ast.Name):
-                            # Check in local symtab first
-                            if key_arg.id in local_sym_tab:
-                                key_var = local_sym_tab[key_arg.id]
-                                key_type = key_var.type.pointee
-                            elif key_arg.id in map_sym_tab:
-                                key_var = map_sym_tab[key_arg.id]
-                                key_type = key_var.type.pointee
-                            else:
-                                print("Key variable "
-                                      f"{key_arg.id} not found in symtabs")
-                                return
-                            print(f"Found key variable {key_arg.id} in symtab")
-                            print(f"Key type: {key_type}")
-                        else:
-                            print("Unsupported lookup key arg")
-                            return
-
-                        # TODO: generate call to bpf_map_lookup_elem
-                        result_ptr = bpf_map_lookup_elem_emitter(
-                            map_global, key_var, module, builder)
-
-                    else:
-                        print(f"Map {map_name} not found in symbol table")
+            if isinstance(rval.func.value, ast.Call) and isinstance(rval.func.value.func, ast.Name):
+                map_name = rval.func.value.func.id
+                method_name = rval.func.attr
+                if map_name in map_sym_tab:
+                    map_ptr = map_sym_tab[map_name]
+                    if method_name in helper_func_list:
+                        val = handle_helper_call(
+                            rval, module, builder, func, local_sym_tab, map_sym_tab)
+                        # var = builder.alloca(ir.IntType(64), name=var_name)
+                        # var.align = 8
+                        builder.store(val, local_sym_tab[var_name])
+                        # local_sym_tab[var_name] = var
             else:
-                print("Unsupported assignment from method call")
+                print("Unsupported assignment call structure")
+        else:
+            print("Unsupported assignment call function type")
+    elif isinstance(rval, ast.BinOp):
+        handle_binary_op(rval, module, builder, func,
+                         local_sym_tab, map_sym_tab)
+    else:
+        print("Unsupported assignment value type")
+
+
+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
+    elif isinstance(cond, ast.Compare):
+        lhs = eval_expr(func, module, builder, cond.left,
+                        local_sym_tab, map_sym_tab)
+        if len(cond.ops) != 1 or len(cond.comparators) != 1:
+            print("Unsupported complex comparison")
+            return None
+        rhs = eval_expr(func, module, builder,
+                        cond.comparators[0], local_sym_tab, map_sym_tab)
+        op = cond.ops[0]
+
+        if lhs.type != rhs.type:
+            if isinstance(lhs.type, ir.IntType) and isinstance(rhs.type, ir.IntType):
+                # Extend the smaller type to the larger type
+                if lhs.type.width < rhs.type.width:
+                    lhs = builder.sext(lhs, rhs.type)
+                elif lhs.type.width > rhs.type.width:
+                    rhs = builder.sext(rhs, lhs.type)
+            else:
+                print("Type mismatch in comparison")
+                return None
+
+        if isinstance(op, ast.Eq):
+            return builder.icmp_signed("==", lhs, rhs)
+        elif isinstance(op, ast.NotEq):
+            return builder.icmp_signed("!=", lhs, rhs)
+        elif isinstance(op, ast.Lt):
+            return builder.icmp_signed("<", lhs, rhs)
+        elif isinstance(op, ast.LtE):
+            return builder.icmp_signed("<=", lhs, rhs)
+        elif isinstance(op, ast.Gt):
+            return builder.icmp_signed(">", lhs, rhs)
+        elif isinstance(op, ast.GtE):
+            return builder.icmp_signed(">=", lhs, rhs)
+        else:
+            print("Unsupported comparison operator")
+            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."""
+    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")
+    if stmt.orelse:
+        else_block = func.append_basic_block(name="if.else")
+    else:
+        else_block = None
+
+    cond = handle_cond(func, module, builder, stmt.test,
+                       local_sym_tab, map_sym_tab)
+    if else_block:
+        builder.cbranch(cond, then_block, else_block)
+    else:
+        builder.cbranch(cond, then_block, merge_block)
+
+    builder.position_at_end(then_block)
+    for s in stmt.body:
+        process_stmt(func, module, builder, s,
+                     local_sym_tab, map_sym_tab, False)
+    if not builder.block.is_terminated:
+        builder.branch(merge_block)
+
+    if else_block:
+        builder.position_at_end(else_block)
+        for s in stmt.orelse:
+            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.AugAssign):
+        raise SyntaxError("Augmented assignment not supported")
+    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):
@@ -118,30 +234,71 @@ def process_func_body(module, builder, func_node, func, ret_type, map_sym_tab):
 
     local_sym_tab = {}
 
+    # pre-allocate dynamic variables
     for stmt in func_node.body:
-        if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Call):
-            call = stmt.value
-            if isinstance(call.func, ast.Name) and call.func.id == "print":
-                bpf_printk_emitter(call, module, builder, func)
-            if isinstance(call.func, ast.Name) and call.func.id == "bpf_ktime_get_ns":
-                bpf_ktime_get_ns_emitter(call, module, builder, func)
-        elif isinstance(stmt, ast.Assign):
-            handle_assign(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}")
+        if isinstance(stmt, ast.Assign):
+            if len(stmt.targets) != 1:
+                print("Unsupported multiassignment")
+                continue
+            target = stmt.targets[0]
+            if not isinstance(target, ast.Name):
+                print("Unsupported assignment target")
+                continue
+            var_name = target.id
+            rval = stmt.value
+            if isinstance(rval, ast.Call):
+                if isinstance(rval.func, ast.Name):
+                    call_type = rval.func.id
+                    if call_type in ("c_int32", "c_int64", "c_uint32", "c_uint64"):
+                        ir_type = ctypes_to_ir(call_type)
+                        var = builder.alloca(ir_type, name=var_name)
+                        var.align = ir_type.width // 8
+                        print(
+                            f"Pre-allocated variable {var_name} of type {call_type}")
+                    elif call_type in helper_func_list:
+                        # Assume return type is int64 for now
+                        ir_type = ir.IntType(64)
+                        var = builder.alloca(ir_type, name=var_name)
+                        var.align = ir_type.width // 8
+                        print(
+                            f"Pre-allocated variable {var_name} for helper")
+                elif isinstance(rval.func, ast.Attribute):
+                    ir_type = ir.PointerType(ir.IntType(64))
+                    var = builder.alloca(ir_type, name=var_name)
+                    # var.align = ir_type.width // 8
+                    print(
+                        f"Pre-allocated variable {var_name} for map")
                 else:
-                    builder.ret(ir.Constant(
-                        ret_type, stmt.value.args[0].value))
-                    did_return = True
+                    print("Unsupported assignment call function type")
+                    continue
+            elif isinstance(rval, ast.Constant):
+                if isinstance(rval.value, bool):
+                    ir_type = ir.IntType(1)
+                    var = builder.alloca(ir_type, name=var_name)
+                    var.align = 1
+                    print(
+                        f"Pre-allocated variable {var_name} of type c_bool")
+                elif isinstance(rval.value, int):
+                    # Assume c_int64 for now
+                    ir_type = ir.IntType(64)
+                    var = builder.alloca(ir_type, name=var_name)
+                    var.align = ir_type.width // 8
+                    print(
+                        f"Pre-allocated variable {var_name} of type c_int64")
+                else:
+                    print("Unsupported constant type")
+                    continue
             else:
-                print("Unsupported return value")
+                print("Unsupported assignment value type")
+                continue
+            local_sym_tab[var_name] = var
+
+    print(f"Local symbol table: {local_sym_tab.keys()}")
+
+    for stmt in func_node.body:
+        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))
 

pythonbpf/helpers.py

@@ -1,4 +1,5 @@
 import ctypes
 
-def bpf_ktime_get_ns():
+
+def ktime():
     return ctypes.c_int64(0)

pythonbpf/maps.py

@@ -16,8 +16,9 @@ class HashMap:
             del self.entries[key]
         else:
             raise KeyError(f"Key {key} not found in map")
-            
-    def update(self, key, value):
+    
+    # TODO: define the flags that can be added
+    def update(self, key, value, flags=None):
         if key in self.entries:
             self.entries[key] = value
         else:

pythonbpf/maps_pass.py

@@ -74,6 +74,14 @@ def create_map_debug_info(module, map_global, map_name, map_params):
         "size": 32,
         "elements": [array_subrange]
     })
+    
+    array_subrange_max_entries = module.add_debug_info("DISubrange", {"count": map_params["max_entries"]})
+    array_type_max_entries = module.add_debug_info("DICompositeType", {
+        "tag": dc.DW_TAG_array_type,
+        "baseType": uint_type,
+        "size": 32,
+        "elements": [array_subrange_max_entries]
+    })
 
     # Create pointer types
     type_ptr = module.add_debug_info("DIDerivedType", {
@@ -84,7 +92,7 @@ def create_map_debug_info(module, map_global, map_name, map_params):
 
     max_entries_ptr = module.add_debug_info("DIDerivedType", {
         "tag": dc.DW_TAG_pointer_type,
-        "baseType": array_type,
+        "baseType": array_type_max_entries,
         "size": 64
     })