csv-cli.py

#!/usr/bin/env python3

"""
A script to output tabular data from a Storm query. Stormfile queries must emit data via $lib.csv.emit()

Based on Synapse's csvtool
    source: https://github.com/vertexproject/synapse/blob/master/synapse/tools/csvtool.py
"""

import argparse
import asyncio
import sys
from typing import Optional

import rich.console
import rich.table
import synapse.telepath as s_telepath

def parse_args() -> Optional[argparse.Namespace]:
    parser = argparse.ArgumentParser(prog="cli-csv", description="Tool to print tabular data from a Synapse cortex")
    parser.add_argument("-c", "--cortex", required=True, help="Cortex telepath/alias")
    parser.add_argument("-n", "--no-header", help="Specify this flag if the provided storm file query doesn't output a header row")

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("-f", "--storm-file", type=str, help="Path to storm file to execute")
    group.add_argument("-q", "--query", type=str, help="Storm query to execute")

    args = parser.parse_args()
    if not args:
        parser.print_help()
        return None
    
    return args

async def main() -> int:
    args = parse_args()
    if not args:
        return 1

    query = ""
    if args.query:
        # if we are running a query, just show the repr for each node
        query = args.query + " | { $lib.csv.emit($node.repr()) }"
    else:
        with open(args.storm_file, "r") as f:
            query = f.read()

        # incoming queries should be doing csv emits
        if not "$lib.csv.emit" in query:
            raise ValueError("Did not detect $lib.csv.emit() calls in the provided query file!")

    async with await s_telepath.openurl(args.cortex) as core:
        opts = {}
        opts["show"] = ("csv:row", "print", "warn", "err")

        rows = []

        async for name, info in core.storm(query, opts=opts):
            if name == "csv:row":
                rows.append(info["row"])
                continue

            if name in ("init", "fini"):
                continue
            
            print(f"{name}: {info}")

    tbl = rich.table.Table(show_lines=True, show_header=(False if args.no_header else True))

    if not args.no_header:
        if args.query:
            tbl.add_column("repr")
        else:
            [tbl.add_column(str(c)) for c in rows[0]]
            rows = rows[1:]

    [tbl.add_row(*[str(c) for c in r]) for r in rows]

    console = rich.console.Console()
    console.print(tbl)

    return 0

if __name__ == "__main__":
    sys.exit(asyncio.run(main()))

example.storm

file:bytes#test.export

+{ -> inet:dns:request +:query:name:fqdn -> inet:fqdn }

$lib.csv.emit("MD5", "FQDN", "A rez")

$row = ({"md5": :md5})

{
    -> inet:dns:request +:query:name:fqdn -> inet:fqdn 
    $row.fqdn = $node.value()

    -> inet:dns:a
    $row.ipv4 = $node.repr(ipv4)
}

$lib.csv.emit($row.md5, $row.fqdn, $row.ipv4)