From 4ce92f0c953011d25b943ac88430850126138d32 Mon Sep 17 00:00:00 2001
From: Pragyansh Chaturvedi <pragyansh.chaturvedi@canonical.com>
Date: Sat, 6 Sep 2025 13:41:56 +0530
Subject: [PATCH] Implement probe string extraction

---
 pythonbpf/codegen.py        |  3 ++-
 pythonbpf/functions_pass.py | 21 +++++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/pythonbpf/codegen.py b/pythonbpf/codegen.py
index b4163be..6fc000d 100644
--- a/pythonbpf/codegen.py
+++ b/pythonbpf/codegen.py
@@ -1,7 +1,7 @@
 import ast
 from llvmlite import ir
 from .license_pass import license_processing
-from .functions_pass import functions_processing
+from .functions_pass import func_proc, functions_processing
 from .constants_pass import constants_processing
 from .globals_pass import globals_processing
 
@@ -26,6 +26,7 @@ def processor(source_code, filename, module):
     for func_node in bpf_chunks:
         print(f"Found BPF function: {func_node.name}")
 
+    func_proc(tree, module, bpf_chunks)
     # For now, we will parse the BPF specific parts of AST
     # Big rewrite
 
diff --git a/pythonbpf/functions_pass.py b/pythonbpf/functions_pass.py
index 3754918..1eb35e8 100644
--- a/pythonbpf/functions_pass.py
+++ b/pythonbpf/functions_pass.py
@@ -43,6 +43,27 @@ def emit_function(module: ir.Module, name: str):
     return func
 
 
+def get_probe_string(func_node):
+    """Extract the probe string from the decorator of the function node."""
+    # TODO: right now we have the whole string in the section decorator
+    # But later we can implement typed tuples for tracepoints and kprobes
+    # For helper functions, we return "helper"
+
+    for decorator in func_node.decorator_list:
+        if isinstance(decorator, ast.Call) and isinstance(decorator.func, ast.Name):
+            if decorator.func.id == "section" and len(decorator.args) == 1:
+                arg = decorator.args[0]
+                if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
+                    return arg.value
+    return "helper"
+
+
+def func_proc(tree, module, chunks):
+    for func_node in chunks:
+        func_type = get_probe_string(func_node)
+        print(f"Found probe_string of {func_node.name}: {func_type}")
+
+
 def functions_processing(tree, module):
     bpf_functions = []
     helper_functions = []