gen_python_md.pl

#!/usr/bin/perl
# This was written by Gemini

use strict;
use warnings;

my $stubs_dir = "../arkime/capture/parsers/stubs";
my $output_file = "_wiki/python.md";

my $header = <<'END_HEADER';
---
donotedit: "DO NOT EDIT MANUALLY - This file is generated by gen_python_md.pl"
title: Python
layout: wiki
permalink: "/python"
copyLink: True
---

<div class="full-height-and-width-container with-footer p-3" markdown="1">

# Python
{: .section-header.mt-1 }

Starting with version 6, Arkime now support Python scripting for custom processing of packets and sessions.
This allows you to write custom classifiers and parsers in Python.
The Python support in Arkime requires Python 3.12 or newer, so it may not be available with older linux distributions.
Use the setting <code>disablePython=true</code> to disable Python support in Arkime.

<strong>Currently Python support is unavailable in the AL2023, Ubuntu 22, and Debian 12 Arkime packages.</strong>
END_HEADER

my $example = <<'END_EXAMPLE';
## Example

Create a <code>/opt/arkime/parsers/example.py</code> file with the following content:

```python
import arkime
import arkime_session
import arkime_packet
import sys

def my_parsers_cb(session, bytes, len, which):
    # Write code here to parse the bytes and extract information
    print("PARSER:", arkime_session.get(session, "ip.src"), ":", arkime_session.get(session, "port.src"), "->", arkime_session.get(session, "ip.dst"), ":", arkime_session.get(session, "port.dst"), "len", len, "which", which)

    # then you could set a field
    # arkime_session.add_string(session, pos, "my value")

    # A parser should return -1 to unregister itself, 0 to continue parsing
    return 0

def my_classify_callback(session, bytes, len, which):
    print("CLASSIFY:", arkime_session.get(session, "ip.src"), ":", arkime_session.get(session, "port.src"), "->", arkime_session.get(session, "ip.dst"), ":", arkime_session.get(session, "port.dst"), "len", len, "which", which)

    # Example of adding a tag
    arkime_session.add_tag(session, "python")

    # Do some kind of check to see if you want to classify this session, if so register
    arkime_session.register_parser(session, my_parsers_cb)

def my_pre_save_callback(session, final):
    print("PRE SAVE:", arkime_session.get(session, "ip.src"), ":", arkime_session.get(session, "port.src"), "->", arkime_session.get(session, "ip.dst"), ":", arkime_session.get(session, "port.dst"), "final", final)

def my_save_callback(session, final):
    print("SAVE:", arkime_session.get(session, "ip.src"), ":", arkime_session.get(session, "port.src"), "->", arkime_session.get(session, "ip.dst"), ":", arkime_session.get(session, "port.dst"), "final", final)

def my_ethernet_cb(batch, packet, bytes, len):
    print("ETHERNET:", "batch", batch, "packet", "packet", "bytes", bytes, "len", len, "pktlen", arkime_packet.get(packet, "pktlen"))

    # Remove first 18 bytes of ethernet header and run ethernet callback again
    bytes = bytes[18:]
    return arkime_packet.run_ethernet_cb(batch, packet, bytes, 0, "example")


### Start ###
# Register a classifier. This example will match all TCP sessions
arkime.register_tcp_classifier("test", 0, bytes("", "ascii"), my_classify_callback)
arkime.register_pre_save(my_pre_save_callback)
arkime.register_save(my_save_callback)

arkime_packet.set_ethernet_cb(0xff12, my_ethernet_cb)

# Create a new field in the session we will be setting
pos = arkime.field_define("arkime_rulz", "kind:lotermfield;db:arkime_rulz")

print("VERSION", arkime.VERSION, "CONFIG_PREFIX", arkime.CONFIG_PREFIX, "POS", pos)
```
</div>
END_EXAMPLE

sub parse_docstring {
    my ($doc) = @_;
    return ("", {{}}) unless defined $doc;

    # Remove leading/trailing quotes and whitespace
    $doc =~ s/^\s*"""\s*//;
    $doc =~ s/\s*"""\s*$//;

    my @lines = split /\n/, $doc;
    my $description = "";
    my %args = ();
    my $current_arg = undef;
    my $in_args = 0;

    foreach my $line (@lines) {
        $line =~ s/^\s*|\u005Cs*$//g; # trim

        if ($line =~ /^Args:/) {
            $in_args = 1;
            next;
        }

        if ($in_args) {
            if ($line =~ /^(\w+):\s*(.*)/) {
                $current_arg = $1;
                $args{$current_arg} = $2;
            } elsif (defined $current_arg && $line ne "") {
                if ($line =~ /^-\s*/) {
                    $args{$current_arg} .= "\n  " . $line;
                } else {
                    $args{$current_arg} .= " " . $line;
                }
            }
        } else {
            $description .= "$line\n";
        }
    }
    chomp $description;
    return ($description, \%args);
}

sub get_file_content {
    my ($file) = @_; 
    my $path = "$stubs_dir/$file";
    open my $fh, '<', $path or die "Could not open $path: $!";
    my $content = do { local $/; <$fh> };
    close $fh;
    # print "Debug: Read " . length($content) . " bytes from $path\n";
    return $content;
}

sub parse_types_callbacks {
    my $content = get_file_content("types.pyi");
    my %callbacks = ();

    # Match: Name = Callable[...] followed by docstring
    # We match Name, then consume until we see """
    while ($content =~ /(\w+)\s*=\s*Callable.*?"""(.*?)"""/gs) {
        my $name = $1;
        my $doc = $2;
        $callbacks{$name} = $doc;
    }
    return %callbacks;
}

my %callbacks_doc = parse_types_callbacks();

open my $out, '>', $output_file or die "Could not open $output_file: $!";
print $out $header;

# === Arkime Module ===
print $out "\n# Python Arkime Module\n";
print $out "\nThe Python Arkime module has high level methods to register callbacks for packet processing.\n";

print $out "\n## Constants\n";
my $arkime_content = get_file_content("arkime.pyi");

# Constants
if ($arkime_content =~ /VERSION: str/) {
    print $out " * VERSION : String - The Arkime version as a string\n";
}
if ($arkime_content =~ /CONFIG_PREFIX: str/) {
    print $out " * CONFIG_PREFIX : String - The Arkime install prefix, usually /opt/arkime\n";
}
if ($arkime_content =~ /API_VERSION: int/) {
    print $out " * API_VERSION : Integer - The Arkime API version from arkime.h\n";
}

# PortKind
print $out "\n### PortKind - The PortKind for register_port_classifier. Bitwise OR the values together to match multiple ports.\n";
while ($arkime_content =~ /^(PORT_\w+):\s*PortKind\s*#\s*(.*)$/gm) {
    my $name = $1;
    my $desc = $2;
    print $out " * $name : Integer - $desc\n";
}

print $out "\n## Callbacks\n";

sub print_callback {
    my ($name, $sig) = @_;
    if (exists $callbacks_doc{$name}) {
        my ($desc, $args_ref) = parse_docstring($callbacks_doc{$name});
        print $out "\n### $sig\n";
        print $out "$desc\n";
        # Re-parse for ordered args
        my $doc = $callbacks_doc{$name};
        $doc =~ s/^\s*"""\s*//; $doc =~ s/\s*"""\s*$//;
        my @lines = split /\n/, $doc;
        my $in_args = 0;
        my $current_arg = undef;
        
        foreach my $line (@lines) {
            $line =~ s/^\s*|\u005Cs*$//g;
            if ($line =~ /^Args:/) { $in_args = 1; next; }
            if ($in_args) {
                 if ($line =~ /^(\w+):\s*(.*)/) {
                     $current_arg = $1;
                     print $out "* $current_arg: $2\n";
                 } elsif (defined $current_arg && $line ne "") {
                     if ($line =~ /^-\s*/) {
                         print $out "  $line\n";
                     } else {
                         print $out " $line\n";
                     }
                 }
            }
        }
    }
}

print_callback("ClassifyCb", "classifyCb(session, packetBytes, packetLen, direction)");
print_callback("ParserCb", "long parserCb(session, packetBytes, packetLen, direction)");
print_callback("SaveCb", "saveCb(session, final)");

print $out "\n## Methods\n";

sub print_functions {

    my ($content) = @_;    

    my @functions = ();

    # Match def function_name ... """docstring"""
    # We match def, then name, then everything until """
        while ($content =~ /def\s+(\w+)(.*?)("""(.*?)""")/gs) {
            my $name = $1;
            my $rest_of_sig = $2;
            my $doc = $4; # Capture group 4 is the content inside the triple quotes
            
            my ($desc, $args_ref) = parse_docstring($doc);
            
            # Get order of args from docstring if possible
            my @doc_args = ();
            my $d_copy = $doc;
            $d_copy =~ s/^\s*"""\s*//; $d_copy =~ s/\s*"""\s*$//;
            my $in_a = 0;
            my $arg_indent = 0;
            foreach my $line (split /\n/, $d_copy) {
                my $orig_line = $line;
                $line =~ s/^\s*|\s*$//g;
                if ($line =~ /^Args:/) { $in_a = 1; next; }
                if ($line =~ /^Returns:/) { $in_a = 0; next; }
                if ($in_a && $line =~ /^(\w+):/) {
                    my $arg_name = $1;
                    # Only treat as new arg if at arg indent level (8 spaces typically)
                    $orig_line =~ /^(\s*)/;
                    my $indent = length($1);
                    if ($arg_indent == 0) { $arg_indent = $indent; }
                    if ($indent <= $arg_indent) {
                        push @doc_args, $arg_name;
                    }
                }
            }
            
            my $sig_args = "";
            if (@doc_args) {
                $sig_args = join(", ", @doc_args);
            } else {
                 # Fallback to extracting from rest_of_sig if no Args section
                 if ($rest_of_sig =~ /\((.*)\)/s) {
                     my $raw_params = $1;
                     # Remove type hints crudely
                     $raw_params =~ s/:\s*[^,)]+//g; 
                     $raw_params =~ s/\s//g;
                     $raw_params =~ s/,/, /g;
                     $sig_args = $raw_params;
                 }
            }
    
        
        my $sig = "$name($sig_args)";
        
        if ($name eq "field_define") { $sig = "fieldPos field_define(fieldExpression, fieldDefinition)"; }
        elsif ($name eq "field_get") { $sig = "fieldPos field_get(fieldExpression)"; }
        
        push @functions, {
            name => $name,
            sig => $sig,
            desc => $desc,
            d_copy => $d_copy
        };
    }

    # Sort by name
    @functions = sort { $a->{name} cmp $b->{name} } @functions;

    foreach my $func (@functions) {
        print $out "\n### $func->{sig}\n";
        print $out "$func->{desc}\n";
        
        # Args and Returns
        my $in_args = 0;
        my $in_returns = 0;
        my $current_arg = undef;
        my $arg_indent = 0;
        foreach my $line (split /\n/, $func->{d_copy}) {
            my $orig_line = $line;
            $line =~ s/^\s*|\s*$//g;
            if ($line =~ /^Args:/) { $in_args = 1; $in_returns = 0; next; }
            if ($line =~ /^Returns:/) { $in_args = 0; $in_returns = 1; $arg_indent = 0; next; }
            if ($in_args) {
                $orig_line =~ /^(\s*)/;
                my $indent = length($1);
                if ($line =~ /^(\w+):\s*(.*)/) {
                    if ($arg_indent == 0) { $arg_indent = $indent; }
                    if ($indent <= $arg_indent) {
                        $current_arg = $1;
                        print $out "* $current_arg: $2\n";
                    } elsif (defined $current_arg) {
                        print $out "  $line\n";
                    }
                } elsif (defined $current_arg && $line ne "") {
                    if ($line =~ /^-\s*/) {
                        print $out "  $line\n";
                    } else {
                        print $out " $line\n";
                    }
                }
            } elsif ($in_returns && $line ne "") {
                $orig_line =~ /^(\s*)/;
                my $indent = length($1);
                if ($line =~ /^(\w+):\s*(.*)/) {
                    if ($arg_indent == 0) { $arg_indent = $indent; }
                    if ($indent <= $arg_indent) {
                        print $out "\nReturns $1: $2\n";
                    } else {
                        print $out " $line\n";
                    }
                } else {
                    print $out " $line\n";
                }
            }
        }
    }
}


print_functions($arkime_content);

print $out "\n## Variables\n";
print $out "\n### CONFIG_PREFIX\n\nThe Arkime install prefix, usually /opt/arkime\n";
print $out "\n### VERSION\n\nThe Arkime version as a string\n";


# === Arkime Session Module ===
print $out "\n# Python Arkime Session Module\n";
my $session_content = get_file_content("arkime_session.pyi");
if ($session_content =~ /"""(.*?)"""/s) {
    my $mod_doc = $1;
    $mod_doc =~ s/^\s*|\u005Cs*$//g;
    # Remove title if duplicated
    $mod_doc =~ s/^Python Arkime Session Module\s*//i;
    print $out "\n$mod_doc\n";
}

print $out "\n## Methods\n";
print_functions($session_content);

# === Arkime Packet Module ===
print $out "\n# Python Arkime Packet Module\n";
my $packet_content = get_file_content("arkime_packet.pyi");
if ($packet_content =~ /"""(.*?)"""/s) {
    my $mod_doc = $1;
    $mod_doc =~ s/^\s*|\u005Cs*$//g;
    # Remove title if duplicated
    $mod_doc =~ s/^Python Arkime Packet Module\s*//i;
    print $out "\n$mod_doc\n";
}

print $out "\n## Constants\n";
print $out "\n### PacketRC\n";
print $out "The return values for a packetCb callback.\n";
while ($packet_content =~ /^(\w+):\s*PacketRC\s*#\s*(.*)$/gm) {
    my $name = $1;
    my $desc = $2;
    print $out " * $name : Integer - $desc\n";
}

print $out "\n## Callbacks\n";
print_callback("PacketCb", "PacketRC packetCb(batch, packet, packetBytes, packetLen)");

print $out "\n## Methods\n";
print_functions($packet_content);

print $out $example;

close $out;
print "Generated $output_file\n";