{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "22dd4e7b-2ea2-49cb-a8d5-1da108c10034",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "from pythonbpf import bpf, map, section, bpfglobal, BPF\n",
    "from pythonbpf.helpers import pid\n",
    "from pythonbpf.maps import HashMap\n",
    "from pylibbpf import *\n",
    "from ctypes import c_void_p, c_int64, c_uint64, c_int32\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ac7a07bf-440f-41e2-bec8-95f520f9cd53",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Module(\n",
      "    body=[\n",
      "        FunctionDef(\n",
      "            name='hist',\n",
      "            args=arguments(\n",
      "                posonlyargs=[],\n",
      "                args=[],\n",
      "                kwonlyargs=[],\n",
      "                kw_defaults=[],\n",
      "                defaults=[]),\n",
      "            body=[\n",
      "                Return(\n",
      "                    value=Call(\n",
      "                        func=Name(id='HashMap', ctx=Load()),\n",
      "                        args=[],\n",
      "                        keywords=[\n",
      "                            keyword(\n",
      "                                arg='key',\n",
      "                                value=Name(id='c_int32', ctx=Load())),\n",
      "                            keyword(\n",
      "                                arg='value',\n",
      "                                value=Name(id='c_uint64', ctx=Load())),\n",
      "                            keyword(\n",
      "                                arg='max_entries',\n",
      "                                value=Constant(value=4096))]))],\n",
      "            decorator_list=[\n",
      "                Name(id='bpf', ctx=Load()),\n",
      "                Name(id='map', ctx=Load())],\n",
      "            returns=Name(id='HashMap', ctx=Load()),\n",
      "            type_params=[]),\n",
      "        FunctionDef(\n",
      "            name='hello',\n",
      "            args=arguments(\n",
      "                posonlyargs=[],\n",
      "                args=[\n",
      "                    arg(\n",
      "                        arg='ctx',\n",
      "                        annotation=Name(id='c_void_p', ctx=Load()))],\n",
      "                kwonlyargs=[],\n",
      "                kw_defaults=[],\n",
      "                defaults=[]),\n",
      "            body=[\n",
      "                Assign(\n",
      "                    targets=[\n",
      "                        Name(id='process_id', ctx=Store())],\n",
      "                    value=Call(\n",
      "                        func=Name(id='pid', ctx=Load()),\n",
      "                        args=[],\n",
      "                        keywords=[])),\n",
      "                Assign(\n",
      "                    targets=[\n",
      "                        Name(id='one', ctx=Store())],\n",
      "                    value=Constant(value=1)),\n",
      "                Assign(\n",
      "                    targets=[\n",
      "                        Name(id='prev', ctx=Store())],\n",
      "                    value=Call(\n",
      "                        func=Attribute(\n",
      "                            value=Call(\n",
      "                                func=Name(id='hist', ctx=Load()),\n",
      "                                args=[],\n",
      "                                keywords=[]),\n",
      "                            attr='lookup',\n",
      "                            ctx=Load()),\n",
      "                        args=[\n",
      "                            Name(id='process_id', ctx=Load())],\n",
      "                        keywords=[])),\n",
      "                If(\n",
      "                    test=Name(id='prev', ctx=Load()),\n",
      "                    body=[\n",
      "                        Assign(\n",
      "                            targets=[\n",
      "                                Name(id='previous_value', ctx=Store())],\n",
      "                            value=BinOp(\n",
      "                                left=Name(id='prev', ctx=Load()),\n",
      "                                op=Add(),\n",
      "                                right=Constant(value=1))),\n",
      "                        Expr(\n",
      "                            value=Call(\n",
      "                                func=Name(id='print', ctx=Load()),\n",
      "                                args=[\n",
      "                                    JoinedStr(\n",
      "                                        values=[\n",
      "                                            Constant(value='count: '),\n",
      "                                            FormattedValue(\n",
      "                                                value=Name(id='previous_value', ctx=Load()),\n",
      "                                                conversion=-1),\n",
      "                                            Constant(value=' with '),\n",
      "                                            FormattedValue(\n",
      "                                                value=Name(id='process_id', ctx=Load()),\n",
      "                                                conversion=-1)])],\n",
      "                                keywords=[])),\n",
      "                        Expr(\n",
      "                            value=Call(\n",
      "                                func=Attribute(\n",
      "                                    value=Call(\n",
      "                                        func=Name(id='hist', ctx=Load()),\n",
      "                                        args=[],\n",
      "                                        keywords=[]),\n",
      "                                    attr='update',\n",
      "                                    ctx=Load()),\n",
      "                                args=[\n",
      "                                    Name(id='process_id', ctx=Load()),\n",
      "                                    Name(id='previous_value', ctx=Load())],\n",
      "                                keywords=[])),\n",
      "                        Return(\n",
      "                            value=Call(\n",
      "                                func=Name(id='c_int64', ctx=Load()),\n",
      "                                args=[\n",
      "                                    Constant(value=0)],\n",
      "                                keywords=[]))],\n",
      "                    orelse=[\n",
      "                        Expr(\n",
      "                            value=Call(\n",
      "                                func=Attribute(\n",
      "                                    value=Call(\n",
      "                                        func=Name(id='hist', ctx=Load()),\n",
      "                                        args=[],\n",
      "                                        keywords=[]),\n",
      "                                    attr='update',\n",
      "                                    ctx=Load()),\n",
      "                                args=[\n",
      "                                    Name(id='process_id', ctx=Load()),\n",
      "                                    Name(id='one', ctx=Load())],\n",
      "                                keywords=[]))]),\n",
      "                Return(\n",
      "                    value=Call(\n",
      "                        func=Name(id='c_int64', ctx=Load()),\n",
      "                        args=[\n",
      "                            Constant(value=0)],\n",
      "                        keywords=[]))],\n",
      "            decorator_list=[\n",
      "                Name(id='bpf', ctx=Load()),\n",
      "                Call(\n",
      "                    func=Name(id='section', ctx=Load()),\n",
      "                    args=[\n",
      "                        Constant(value='tracepoint/syscalls/sys_enter_clone')],\n",
      "                    keywords=[])],\n",
      "            returns=Name(id='c_int64', ctx=Load()),\n",
      "            type_params=[]),\n",
      "        FunctionDef(\n",
      "            name='LICENSE',\n",
      "            args=arguments(\n",
      "                posonlyargs=[],\n",
      "                args=[],\n",
      "                kwonlyargs=[],\n",
      "                kw_defaults=[],\n",
      "                defaults=[]),\n",
      "            body=[\n",
      "                Return(\n",
      "                    value=Constant(value='GPL'))],\n",
      "            decorator_list=[\n",
      "                Name(id='bpf', ctx=Load()),\n",
      "                Name(id='bpfglobal', ctx=Load())],\n",
      "            returns=Name(id='str', ctx=Load()),\n",
      "            type_params=[]),\n",
      "        Assign(\n",
      "            targets=[\n",
      "                Name(id='b', ctx=Store())],\n",
      "            value=Call(\n",
      "                func=Name(id='BPF', ctx=Load()),\n",
      "                args=[],\n",
      "                keywords=[]))],\n",
      "    type_ignores=[])\n",
      "Found BPF function/struct: hist\n",
      "Found BPF function/struct: hello\n",
      "Found BPF function/struct: LICENSE\n",
      "Found BPF map: hist\n",
      "Processing BPF map: hist\n",
      "Creating HashMap map: hist\n",
      "Map parameters: {'type': 'HASH', 'key': 'c_int32', 'value': 'c_uint64', 'max_entries': 4096}\n",
      "Created BPF map: hist\n",
      "Found probe_string of hello: tracepoint/syscalls/sys_enter_clone\n",
      "Pre-allocated variable process_id for helper\n",
      "Pre-allocated variable one of type c_int64\n",
      "Pre-allocated variable prev for map\n",
      "Pre-allocated variable previous_value of type c_int64\n",
      "Local symbol table: dict_keys(['process_id', 'one', 'prev', 'previous_value'])\n",
      "Processing statement: Assign(targets=[Name(id='process_id', ctx=Store())], value=Call(func=Name(id='pid', ctx=Load()), args=[], keywords=[]))\n",
      "Handling assignment to Name(id='process_id', ctx=Store())\n",
      "Assignment call type: pid\n",
      "Assigned constant pid to process_id\n",
      "Processing statement: Assign(targets=[Name(id='one', ctx=Store())], value=Constant(value=1))\n",
      "Handling assignment to Name(id='one', ctx=Store())\n",
      "Assigned constant 1 to one\n",
      "Processing statement: Assign(targets=[Name(id='prev', ctx=Store())], value=Call(func=Attribute(value=Call(func=Name(id='hist', ctx=Load()), args=[], keywords=[]), attr='lookup', ctx=Load()), args=[Name(id='process_id', ctx=Load())], keywords=[]))\n",
      "Handling assignment to Name(id='prev', ctx=Store())\n",
      "Assignment call attribute: Attribute(value=Call(func=Name(id='hist', ctx=Load()), args=[], keywords=[]), attr='lookup', ctx=Load())\n",
      "{}\n",
      "Processing statement: If(test=Name(id='prev', ctx=Load()), body=[Assign(targets=[Name(id='previous_value', ctx=Store())], value=BinOp(left=Name(id='prev', ctx=Load()), op=Add(), right=Constant(value=1))), Expr(value=Call(func=Name(id='print', ctx=Load()), args=[JoinedStr(values=[Constant(value='count: '), FormattedValue(value=Name(id='previous_value', ctx=Load()), conversion=-1), Constant(value=' with '), FormattedValue(value=Name(id='process_id', ctx=Load()), conversion=-1)])], keywords=[])), Expr(value=Call(func=Attribute(value=Call(func=Name(id='hist', ctx=Load()), args=[], keywords=[]), attr='update', ctx=Load()), args=[Name(id='process_id', ctx=Load()), Name(id='previous_value', ctx=Load())], keywords=[])), Return(value=Call(func=Name(id='c_int64', ctx=Load()), args=[Constant(value=0)], keywords=[]))], orelse=[Expr(value=Call(func=Attribute(value=Call(func=Name(id='hist', ctx=Load()), args=[], keywords=[]), attr='update', ctx=Load()), args=[Name(id='process_id', ctx=Load()), Name(id='one', ctx=Load())], keywords=[]))])\n",
      "Handling if statement\n",
      "Processing statement: Assign(targets=[Name(id='previous_value', ctx=Store())], value=BinOp(left=Name(id='prev', ctx=Load()), op=Add(), right=Constant(value=1)))\n",
      "Handling assignment to Name(id='previous_value', ctx=Store())\n",
      "; ModuleID = \"/tmp/tmp98sm2v54.py\"\n",
      "target triple = \"bpf\"\n",
      "target datalayout = \"e-m:e-p:64:64-i64:64-i128:128-n32:64-S128\"\n",
      "\n",
      "@\"hist\" = dso_local global {ptr, ptr, ptr, ptr} zeroinitializer, section \".maps\", align 8, !dbg !22\n",
      "define dso_local i64 @\"hello\"(ptr nocapture %\".1\") noinline nounwind optnone section \"tracepoint/syscalls/sys_enter_clone\"\n",
      "{\n",
      "entry:\n",
      "  %\"process_id\" = alloca i64, align 8\n",
      "  %\"one\" = alloca i64, align 8\n",
      "  %\"prev\" = alloca i64*\n",
      "  %\"previous_value\" = alloca i64, align 8\n",
      "  %\".3\" = inttoptr i64 14 to i64 ()*\n",
      "  %\".4\" = call i64 %\".3\"()\n",
      "  %\".5\" = and i64 %\".4\", 4294967295\n",
      "  store i64 %\".5\", i64* %\"process_id\"\n",
      "  store i64 1, i64* %\"one\"\n",
      "  %\".8\" = inttoptr i64 1 to ptr (ptr, ptr)*\n",
      "  %\".9\" = call ptr %\".8\"({ptr, ptr, ptr, ptr}* @\"hist\", i64* %\"process_id\")\n",
      "  store ptr %\".9\", i64** %\"prev\"\n",
      "  %\".11\" = load i64*, i64** %\"prev\"\n",
      "  %\".12\" = icmp ne i64* %\".11\", null\n",
      "  br i1 %\".12\", label %\"if.then\", label %\"if.else\"\n",
      "if.then:\n",
      "if.end:\n",
      "if.else:\n",
      "}\n",
      "\n",
      "!llvm.dbg.cu = !{ !1 }\n",
      "!0 = !DIFile(directory: \"/tmp\", filename: \"/tmp/tmp98sm2v54.py\")\n",
      "!1 = distinct !DICompileUnit(emissionKind: 1, file: !0, isOptimized: true, language: 29, nameTableKind: 0, producer: \"PythonBPF DSL Compiler\", runtimeVersion: 0, splitDebugInlining: false)\n",
      "!2 = !DIBasicType(encoding: 7, name: \"unsigned int\", size: 32)\n",
      "!3 = !DIBasicType(encoding: 7, name: \"unsigned long long\", size: 64)\n",
      "!4 = !DISubrange(count: 1)\n",
      "!5 = !{ !4 }\n",
      "!6 = !DICompositeType(baseType: !2, elements: !5, size: 32, tag: 1)\n",
      "!7 = !DIDerivedType(baseType: !6, size: 64, tag: 15)\n",
      "!8 = !DIDerivedType(baseType: !2, size: 64, tag: 15)\n",
      "!9 = !DIDerivedType(baseType: !3, size: 64, tag: 15)\n",
      "!10 = !DIDerivedType(baseType: !7, file: !0, name: \"type\", offset: 0, size: 64, tag: 13)\n",
      "!11 = !DIDerivedType(baseType: !8, file: !0, name: \"key\", offset: 64, size: 64, tag: 13)\n",
      "!12 = !DIDerivedType(baseType: !9, file: !0, name: \"value\", offset: 128, size: 64, tag: 13)\n",
      "!13 = !DISubrange(count: 4096)\n",
      "!14 = !{ !13 }\n",
      "!15 = !DICompositeType(baseType: !2, elements: !14, size: 32, tag: 1)\n",
      "!16 = !DIDerivedType(baseType: !15, size: 64, tag: 15)\n",
      "!17 = !DIDerivedType(baseType: !16, file: !0, name: \"max_entries\", offset: 192, size: 64, tag: 13)\n",
      "!18 = !{ !10, !11, !12, !17 }\n",
      "!19 = distinct !DICompositeType(elements: !18, file: !0, size: 256, tag: 19)\n",
      "!20 = distinct !DIGlobalVariable(file: !0, isDefinition: true, isLocal: false, name: \"hist\", scope: !1, type: !19)\n",
      "!21 = !DIExpression()\n",
      "!22 = !DIGlobalVariableExpression(expr: !21, var: !20)\n",
      "left is %\".15\" = load i64, i64* %\".14\", right is i64 1, op is <ast.Add object at 0x70244bec1410>\n",
      "Processing statement: Expr(value=Call(func=Name(id='print', ctx=Load()), args=[JoinedStr(values=[Constant(value='count: '), FormattedValue(value=Name(id='previous_value', ctx=Load()), conversion=-1), Constant(value=' with '), FormattedValue(value=Name(id='process_id', ctx=Load()), conversion=-1)])], keywords=[]))\n",
      "Handling expression: Expr(value=Call(func=Name(id='print', ctx=Load()), args=[JoinedStr(values=[Constant(value='count: '), FormattedValue(value=Name(id='previous_value', ctx=Load()), conversion=-1), Constant(value=' with '), FormattedValue(value=Name(id='process_id', ctx=Load()), conversion=-1)])], keywords=[]))\n",
      "Evaluating expression: Call(func=Name(id='print', ctx=Load()), args=[JoinedStr(values=[Constant(value='count: '), FormattedValue(value=Name(id='previous_value', ctx=Load()), conversion=-1), Constant(value=' with '), FormattedValue(value=Name(id='process_id', ctx=Load()), conversion=-1)])], keywords=[])\n",
      "Value in f-string: Constant(value='count: ')\n",
      "Value in f-string: FormattedValue(value=Name(id='previous_value', ctx=Load()), conversion=-1)\n",
      "Formatted value: FormattedValue(value=Name(id='previous_value', ctx=Load()), conversion=-1)\n",
      "Formatted value variable: %\"previous_value\" = alloca i64, align 8 i64\n",
      "Value in f-string: Constant(value=' with ')\n",
      "Value in f-string: FormattedValue(value=Name(id='process_id', ctx=Load()), conversion=-1)\n",
      "Formatted value: FormattedValue(value=Name(id='process_id', ctx=Load()), conversion=-1)\n",
      "Formatted value variable: %\"process_id\" = alloca i64, align 8 i64\n",
      "Evaluating expression: Name(id='previous_value', ctx=Load())\n",
      "Evaluating expression: Name(id='process_id', ctx=Load())\n",
      "Processing statement: Expr(value=Call(func=Attribute(value=Call(func=Name(id='hist', ctx=Load()), args=[], keywords=[]), attr='update', ctx=Load()), args=[Name(id='process_id', ctx=Load()), Name(id='previous_value', ctx=Load())], keywords=[]))\n",
      "Handling expression: Expr(value=Call(func=Attribute(value=Call(func=Name(id='hist', ctx=Load()), args=[], keywords=[]), attr='update', ctx=Load()), args=[Name(id='process_id', ctx=Load()), Name(id='previous_value', ctx=Load())], keywords=[]))\n",
      "Evaluating expression: Call(func=Attribute(value=Call(func=Name(id='hist', ctx=Load()), args=[], keywords=[]), attr='update', ctx=Load()), args=[Name(id='process_id', ctx=Load()), Name(id='previous_value', ctx=Load())], keywords=[])\n",
      "Handling method call: Attribute(value=Call(func=Name(id='hist', ctx=Load()), args=[], keywords=[]), attr='update', ctx=Load())\n",
      "{}\n",
      "Processing statement: Return(value=Call(func=Name(id='c_int64', ctx=Load()), args=[Constant(value=0)], keywords=[]))\n",
      "Processing statement: Expr(value=Call(func=Attribute(value=Call(func=Name(id='hist', ctx=Load()), args=[], keywords=[]), attr='update', ctx=Load()), args=[Name(id='process_id', ctx=Load()), Name(id='one', ctx=Load())], keywords=[]))\n",
      "Handling expression: Expr(value=Call(func=Attribute(value=Call(func=Name(id='hist', ctx=Load()), args=[], keywords=[]), attr='update', ctx=Load()), args=[Name(id='process_id', ctx=Load()), Name(id='one', ctx=Load())], keywords=[]))\n",
      "Evaluating expression: Call(func=Attribute(value=Call(func=Name(id='hist', ctx=Load()), args=[], keywords=[]), attr='update', ctx=Load()), args=[Name(id='process_id', ctx=Load()), Name(id='one', ctx=Load())], keywords=[])\n",
      "Handling method call: Attribute(value=Call(func=Name(id='hist', ctx=Load()), args=[], keywords=[]), attr='update', ctx=Load())\n",
      "{}\n",
      "Processing statement: Return(value=Call(func=Name(id='c_int64', ctx=Load()), args=[Constant(value=0)], keywords=[]))\n",
      "IR written to /tmp/tmpx05plfu1.ll\n"
     ]
    }
   ],
   "source": [
    "@bpf\n",
    "@map\n",
    "def hist() -> HashMap:\n",
    "    return HashMap(key=c_int32, value=c_uint64, max_entries=4096)\n",
    "\n",
    "@bpf\n",
    "@section(\"tracepoint/syscalls/sys_enter_clone\")\n",
    "def hello(ctx: c_void_p) -> c_int64:\n",
    "    process_id = pid()\n",
    "    one = 1\n",
    "    prev = hist().lookup(process_id)\n",
    "    if prev:\n",
    "        previous_value = prev + 1\n",
    "        print(f\"count: {previous_value} with {process_id}\")\n",
    "        hist().update(process_id, previous_value)\n",
    "        return c_int64(0)\n",
    "    else:\n",
    "        hist().update(process_id, one)\n",
    "    return c_int64(0)\n",
    "\n",
    "\n",
    "@bpf\n",
    "@bpfglobal\n",
    "def LICENSE() -> str:\n",
    "    return \"GPL\"\n",
    "\n",
    "b = BPF()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "93fae9f8-464e-48d6-b61e-57b9f93e508a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recording\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 20376 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "b.load_and_attach()\n",
    "hist = BpfMap(b, hist)\n",
    "print(\"Recording\")\n",
    "time.sleep(10)\n",
    "\n",
    "counts = list(hist.values())\n",
    "\n",
    "plt.hist(counts, bins=10)\n",
    "plt.xlabel(\"Clone calls per PID\")\n",
    "plt.ylabel(\"Frequency\")\n",
    "plt.title(\"Syscall clone counts\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
