simple RST pre-processor to embed storm queries

synapse/tests/test_tools_rstorm.py

@@ -0,0 +1,89 @@
+import synapse.exc as s_exc
+import synapse.tests.utils as s_test
+
+import synapse.common as s_common
+import synapse.tools.rstorm as s_rstorm
+
+rst_in = '''
+HI
+##
+.. storm-cortex:: synapse.tools.rstorm.cortex
+.. storm-cortex:: synapse.tools.rstorm.cortex
+.. storm-opts:: {"vars": {"foo": 10, "bar": "baz"}}
+.. storm-pre:: [ inet:asn=$foo ]
+.. storm:: $lib.print($bar) $lib.warn(omgomgomg)
+.. storm-expect:: baz
+'''
+
+rst_out = '''
+HI
+##
+::
+
+    > $lib.print($bar) $lib.warn(omgomgomg)
+    baz
+    WARNING: omgomgomg
+
+'''
+
+boom1 = '''
+
+.. storm:: $lib.print(newp)
+
+'''
+
+boom2 = '''
+
+.. storm-pre:: $lib.print(newp)
+
+'''
+
+boom3 = '''
+
+.. storm-cortex:: synapse.tools.rstorm.cortex
+.. storm:: $x = (10 + "foo")
+
+'''
+
+class RStormTest(s_test.SynTest):
+
+    async def test_rstorm(self):
+
+        with self.getTestDir() as dirn:
+
+            path = s_common.genpath(dirn, 'test.rst')
+            with s_common.genfile(path) as fd:
+                fd.write(rst_in.encode())
+
+            outpath = s_common.genpath(dirn, 'out.rst')
+
+            await s_rstorm.main(('--save', outpath, path))
+
+            with s_common.genfile(outpath) as fd:
+                text = fd.read().decode()
+
+            self.eq(text, rst_out)
+
+            # boom1 test
+            path = s_common.genpath(dirn, 'boom1.rst')
+            with s_common.genfile(path) as fd:
+                fd.write(boom1.encode())
+
+            with self.raises(s_exc.NoSuchVar):
+                await s_rstorm.main(('--save', outpath, path))
+
+            # boom2 test
+            path = s_common.genpath(dirn, 'boom2.rst')
+            with s_common.genfile(path) as fd:
+                fd.write(boom2.encode())
+
+            with self.raises(s_exc.NoSuchVar):
+                await s_rstorm.main(('--save', outpath, path))
+
+            # boom3 test
+            path = s_common.genpath(dirn, 'boom3.rst')
+            with s_common.genfile(path) as fd:
+                fd.write(boom3.encode())
+
+            with self.raises(s_exc.StormRuntimeError):
+                await s_rstorm.main(('--save', outpath, path))

synapse/tools/rstorm.py

@@ -0,0 +1,134 @@
+import sys
+import json
+import asyncio
+import argparse
+
+import synapse.exc as s_exc
+import synapse.common as s_common
+import synapse.cortex as s_cortex
+
+import synapse.lib.base as s_base
+import synapse.lib.output as s_output
+import synapse.lib.dyndeps as s_dyndeps
+
+async def cortex():
+    base = await s_base.Base.anit()
+    dirn = await base.enter_context(s_common.getTempDir())
+    core = await s_cortex.Cortex.anit(dirn)
+    core.onfini(base.fini)
+    return core
+
+class StormOutput:
+    '''
+    Produce standard output from a stream of storm runtime messages.
+    '''
+    # TODO: Eventually make this obey all cmdr options and swap it in
+
+async def genStormRst(path, debug=False):
+
+    outp = []
+    context = {}
+
+    with open(path, 'r') as fd:
+        lines = fd.readlines()
+
+    for line in lines:
+
+        if line.startswith('.. storm-cortex::'):
+            ctor = line.split('::', 1)[1].strip()
+            core = await (s_dyndeps.getDynLocal(ctor))()
+            if context.get('cortex') is not None:
+                await (context.pop('cortex')).fini()
+            context['cortex'] = core
+            continue
+
+        if line.startswith('.. storm-opts::'):
+            item = json.loads(line.split('::', 1)[1].strip())
+            context['opts'] = item
+            continue
+
+        if line.startswith('.. storm-expect::'):
+            # TODO handle some light weight output confirmation.
+            continue
+
+        if line.startswith('.. storm-pre::'):
+            # runt a storm query to prepare the cortex (do not output)
+
+            text = line.split('::', 1)[1].strip()
+
+            core = context.get('cortex')
+            if core is None:
+                mesg = 'No cortex set.  Use .. storm-cortex::'
+                raise s_exc.NoSuchVar(mesg=mesg)
+
+            opts = context.get('opts')
+            await core.callStorm(text, opts=opts)
+            continue
+
+        if line.startswith('.. storm::'):
+
+            text = line.split('::', 1)[1].strip()
+
+            core = context.get('cortex')
+            if core is None:
+                mesg = 'No cortex set.  Use .. storm-cortex::'
+                raise s_exc.NoSuchVar(mesg=mesg)
+
+            outp.append('::\n')
+            outp.append('\n')
+
+            outp.append(f'    > {text}\n')
+
+            opts = context.get('opts')
+            msgs = await core.stormlist(text, opts=opts)
+
+            # TODO use StormOutput
+            for mesg in await core.stormlist(text, opts=opts):
+
+                if mesg[0] == 'print':
+                    ptxt = mesg[1]['mesg']
+                    outp.append(f'    {ptxt}\n')
+                    continue
+
+                if mesg[0] == 'warn':
+                    ptxt = mesg[1]['mesg']
+                    outp.append(f'    WARNING: {ptxt}\n')
+                    continue
+
+                if mesg[0] == 'err':
+                    raise s_exc.StormRuntimeError(mesg=mesg)
+
+            outp.append('\n')
+            continue
+
+        outp.append(line)
+
+    core = context.get('cortex')
+    if core is not None:
+        await core.fini()
+
+    return outp
+
+prog = 'synapse.tools.rstorm'
+descr = 'An RST pre-processor that allows you to embed storm directives.'
+
+async def main(argv, outp=s_output.stdout):
+
+    pars = argparse.ArgumentParser(prog=prog, description=descr)
+    pars.add_argument('rstfile', help='Input RST file with storm directives.')
+    pars.add_argument('--save', help='Output file to save (default: stdout)')
+
+    opts = pars.parse_args(argv)
+    path = s_common.genpath(opts.rstfile)
+
+    lines = await genStormRst(path)
+
+    if opts.save:
+        with open(s_common.genpath(opts.save), 'w') as fd:
+            [fd.write(line) for line in lines]
+    else:
+        for line in lines:
+            outp.printf(line, addnl=False)
+
+if __name__ == '__main__':
+    sys.exit(asyncio.run(main(sys.argv[1:])))