synapse 2.212.0__py311-none-any.whl → 2.214.0__py311-none-any.whl

synapse/tests/test_tools_shutdown.py

@@ -0,0 +1,95 @@
+import synapse.common as s_common
+
+import synapse.lib.cell as s_cell
+
+
+import synapse.tests.utils as s_test
+import synapse.tools.shutdown as s_t_shutdown
+
+class ShutdownToolTest(s_test.SynTest):
+
+    async def test_tool_shutdown_base(self):
+
+        async with self.getTestCore() as core:
+
+            msgs = await core.stormlist('background { $lib.time.sleep(10) }')
+            self.stormHasNoWarnErr(msgs)
+
+            # add a dmon to ensure task.background=True works correctly
+            await core.addStormDmon({
+                'iden': s_common.guid(),
+                'storm': 'while (true) { $lib.time.sleep(1) }',
+            })
+
+            argv = ['--url', core.getLocalUrl(), '--timeout', '0']
+
+            self.eq(1, await s_t_shutdown.main(argv))
+
+            for task in core.boss.ps():
+                if task.name == 'storm':
+                    await task.kill()
+
+            self.eq(0, await s_t_shutdown.main(['--url', core.getLocalUrl()]))
+
+            self.true(await core.waitfini(timeout=1))
+
+        outp = self.getTestOutp()
+        self.eq(1, await s_t_shutdown.main(['--url', 'newp://hehe'], outp=outp))
+        outp.expect('Error while attempting graceful shutdown')
+
+    async def test_tool_shutdown_leader(self):
+
+        async with self.getTestAha() as aha:
+
+            with self.getTestDir() as dirn:
+
+                dirn00 = s_common.genpath(dirn, '00.cell')
+                dirn01 = s_common.genpath(dirn, '01.cell')
+
+                cell00 = await aha.enter_context(self.addSvcToAha(aha, '00.cell', s_cell.Cell, dirn=dirn00))
+                cell01 = await aha.enter_context(self.addSvcToAha(aha, '01.cell', s_cell.Cell, dirn=dirn01,
+                                                                   provinfo={'mirror': 'cell'}))
+                self.true(cell00.isactive)
+                self.false(cell01.isactive)
+
+                await cell01.sync()
+
+                # confirm that graceful shutdown with peers also demotes...
+                outp = self.getTestOutp()
+                argv = ['--url', cell00.getLocalUrl(), '--timeout', '12']
+                self.eq(0, await s_t_shutdown.main(argv, outp=outp))
+
+                self.false(cell00.isactive)
+                self.true(cell01.isactive)
+                self.true(await cell00.waitfini(timeout=12))
+
+                # and that graceful shutdown without any cluster peers works too...
+                outp.clear()
+                argv = ['--url', cell01.getLocalUrl(), '--timeout', '12']
+                self.eq(0, await s_t_shutdown.main(argv, outp=outp))
+                self.true(await cell01.waitfini(timeout=12))
+
+    async def test_tool_shutdown_no_features(self):
+
+        async with self.getTestAha() as aha:
+            aha.features.pop('getAhaSvcsByIden')
+
+            with self.getTestDir() as dirn:
+
+                dirn00 = s_common.genpath(dirn, '00.cell')
+                dirn01 = s_common.genpath(dirn, '01.cell')
+
+                cell00 = await aha.enter_context(self.addSvcToAha(aha, '00.cell', s_cell.Cell, dirn=dirn00))
+                cell01 = await aha.enter_context(self.addSvcToAha(aha, '01.cell', s_cell.Cell, dirn=dirn01,
+                                                                   provinfo={'mirror': 'cell'}))
+                self.true(cell00.isactive)
+                self.false(cell01.isactive)
+
+                await cell01.sync()
+
+                # confirm that graceful shutdown with peers also demotes...
+                outp = self.getTestOutp()
+                argv = ['--url', cell00.getLocalUrl(), '--timeout', '12']
+                with self.getAsyncLoggerStream('synapse.daemon') as stream:
+                    self.eq(1, await s_t_shutdown.main(argv, outp=outp))
+                stream.expect('AHA server does not support feature: getAhaSvcsByIden >= 1')

synapse/tests/test_utils.py

@@ -3,6 +3,7 @@ import time
 import logging
 import unittest
 
+import synapse.exc as s_exc
 import synapse.common as s_common
 
 import synapse.lib.base as s_base
@@ -100,13 +101,33 @@ class TestUtils(s_t_utils.SynTest):
             self.skipIfNoPath('newpDoesNotExist', mesg='hehe')
         self.isin('newpDoesNotExist mesg=hehe', str(cm.exception))
 
-    def test_syntest_logstream(self):
+    async def test_syntest_logstream(self):
         with self.getLoggerStream('synapse.tests.test_utils') as stream:
             logger.error('ruh roh i am a error message')
+
+        stream.expect('error message')
+        with self.raises(s_exc.SynErr):
+            stream.expect('does not exist')
+
+        self.eq(str(stream), 'ruh roh i am a error message\n')
+        self.true(repr(stream).endswith('valu: ruh roh i am a error message>'))
+        self.true(repr(stream).startswith('<synapse.tests.utils.StreamEvent'))
+
         stream.seek(0)
         mesgs = stream.read()
         self.isin('ruh roh', mesgs)
 
+        with self.getAsyncLoggerStream('synapse.tests.test_utils') as stream:
+            logger.error('ruh roh i am a new message')
+
+        stream.expect('new message')
+        with self.raises(s_exc.SynErr):
+            stream.expect('does not exist')
+
+        self.eq(str(stream), 'ruh roh i am a new message\n')
+        self.true(repr(stream).endswith('valu: ruh roh i am a new message>'))
+        self.true(repr(stream).startswith('<synapse.tests.utils.AsyncStreamEvent'))
+
     def test_syntest_logstream_event(self):
 
         @s_common.firethread

synapse/tests/utils.py

@@ -788,14 +788,9 @@ class CmdGenerator:
 
         return retn
 
-class StreamEvent(io.StringIO, threading.Event):
-    '''
-    A combination of a io.StringIO object and a threading.Event object.
-    '''
+class _StreamIOMixin(io.StringIO):
     def __init__(self, *args, **kwargs):
         io.StringIO.__init__(self, *args, **kwargs)
-        threading.Event.__init__(self)
-        self.mesg = ''
 
     def setMesg(self, mesg):
         '''
@@ -810,6 +805,9 @@ class StreamEvent(io.StringIO, threading.Event):
         self.mesg = mesg
         self.clear()
 
+    def __str__(self):
+        return self.getvalue()
+
     def write(self, s):
         io.StringIO.write(self, s)
         if self.mesg and self.mesg in s:
@@ -819,42 +817,59 @@ class StreamEvent(io.StringIO, threading.Event):
         '''Get the messages as jsonlines. May throw Json errors if the captured stream is not jsonlines.'''
         return jsonlines(self.getvalue())
 
-class AsyncStreamEvent(io.StringIO, asyncio.Event):
+    def expect(self, substr: str):
+        '''
+        Check if a string is present in the messages captured by StreamEvent.
+
+        Args:
+            substr (str): String to check for the existence of.
+        '''
+        valu = self.getvalue()
+        if valu.find(substr) == -1:
+            mesg = '%s.expect(%s) not in %s' % (self.__class__.__name__, substr, valu)
+            raise s_exc.SynErr(mesg=mesg)
+
+class StreamEvent(_StreamIOMixin, threading.Event):
     '''
-    A combination of a io.StringIO object and an asyncio.Event object.
+    A combination of a io.StringIO object and a threading.Event object.
     '''
     def __init__(self, *args, **kwargs):
-        io.StringIO.__init__(self, *args, **kwargs)
-        asyncio.Event.__init__(self)
+        _StreamIOMixin.__init__(self, *args, **kwargs)
+        threading.Event.__init__(self)
         self.mesg = ''
 
-    def setMesg(self, mesg):
-        '''
-        Clear the internal event and set a new message that is used to set the event.
-
-        Args:
-            mesg (str): The string to monitor for.
+    def __repr__(self):
+        cls = self.__class__
+        status = 'set' if self._flag else 'unset'
+        if valu := str(self):
+            valu = s_common.trimText(valu).strip()
+            status = f'{status}, valu: {valu}'
+        return f"<{cls.__module__}.{cls.__qualname__} at {id(self):#x}: {status}>"
 
-        Returns:
-            None
-        '''
-        self.mesg = mesg
-        self.clear()
+class AsyncStreamEvent(_StreamIOMixin, asyncio.Event):
+    '''
+    A combination of a io.StringIO object and an asyncio.Event object.
+    '''
+    def __init__(self, *args, **kwargs):
+        _StreamIOMixin.__init__(self, *args, **kwargs)
+        asyncio.Event.__init__(self)
+        self.mesg = ''
 
-    def write(self, s):
-        io.StringIO.write(self, s)
-        if self.mesg and self.mesg in s:
-            self.set()
+    def __repr__(self):
+        cls = self.__class__
+        status = 'set' if self._value else 'unset'
+        if self._waiters:
+            status = f'{status}, waiters:{len(self._waiters)}'
+        if valu := str(self):
+            valu = s_common.trimText(valu).strip()
+            status = f'{status}, valu: {valu}'
+        return f"<{cls.__module__}.{cls.__qualname__} at {id(self):#x}: {status}>"
 
     async def wait(self, timeout=None):
         if timeout is None:
             return await asyncio.Event.wait(self)
         return await s_coro.event_wait(self, timeout=timeout)
 
-    def jsonlines(self) -> typing.List[dict]:
-        '''Get the messages as jsonlines. May throw Json errors if the captured stream is not jsonlines.'''
-        return jsonlines(self.getvalue())
-
 class HttpReflector(s_httpapi.Handler):
     '''Test handler which reflects get/post data back to the caller'''
 

synapse/tools/aha/easycert.py

@@ -45,6 +45,7 @@ async def main(argv, outp=s_output.stdout):
                 cert = c_x509.load_pem_x509_certificate(certbyts.encode())
                 path = cdir._saveCertTo(cert, 'hosts', f'{name}.crt')
                 outp.printf(f'crt saved: {path}')
+                cdir.delHostCsr(name, outp=outp)
                 return 0
             else:
                 csr = cdir.genUserCsr(name, outp=outp)
@@ -52,6 +53,7 @@ async def main(argv, outp=s_output.stdout):
                 cert = c_x509.load_pem_x509_certificate(certbyts.encode())
                 path = cdir._saveCertTo(cert, 'users', f'{name}.crt')
                 outp.printf(f'crt saved: {path}')
+                cdir.delUserCsr(name, outp=outp)
                 return 0
 
 def getArgParser():

synapse/tools/aha/enroll.py

@@ -98,6 +98,9 @@ async def main(argv, outp=s_output.stdout):
                     teleyaml['aha:servers'] = servers
                     s_common.yamlsave(teleyaml, yamlpath)
 
+            # cleanup CSR
+            certdir.delUserCsr(username)
+
     return 0
 
 async def _main(argv, outp=s_output.stdout):  # pragma: no cover

synapse/tools/axon/dump.py

@@ -0,0 +1,155 @@
+import os
+import tarfile
+import tempfile
+
+import synapse.exc as s_exc
+import synapse.common as s_common
+import synapse.telepath as s_telepath
+
+import synapse.lib.cmd as s_cmd
+import synapse.lib.const as s_const
+import synapse.lib.output as s_output
+
+descr = '''
+Dump blobs from a Synapse Axon.
+'''
+
+MAX_SPOOL_SIZE = s_const.mebibyte * 512
+DEFAULT_ROTATE_SIZE = s_const.gigabyte * 4
+
+def getTarName(celliden, start, end):
+    return f'{celliden}.{start:012d}-{end:012d}.tar.gz'
+
+async def dumpBlobs(opts, outp):
+
+    try:
+        async with await s_telepath.openurl(opts.url) as axon:
+
+            cellinfo = await axon.getCellInfo()
+            celltype = cellinfo['cell']['type']
+            if "axon" not in celltype.lower():
+                mesg = f'Axon dump tool only works on axons, not {celltype}'
+                raise s_exc.TypeMismatch(mesg=mesg)
+            celliden = cellinfo['cell']['iden']
+
+            if os.path.exists(opts.outdir) and not os.path.isdir(opts.outdir):
+                raise s_exc.BadDataValu(mesg=f'Specified output directory {opts.outdir} exists but is not a directory.')
+            os.makedirs(opts.outdir, exist_ok=True)
+
+            if opts.statefile is None:
+                statefile_path = os.path.join(opts.outdir, f'{celliden}.yaml')
+            else:
+                statefile_path = opts.statefile
+                if os.path.isdir(statefile_path):
+                    statefile_path = os.path.join(statefile_path, f'{celliden}.yaml')
+            state = {}
+            if os.path.isfile(statefile_path):
+                if (data := s_common.yamlload(statefile_path)) is not None:
+                    state = data
+            opts.statefile = statefile_path
+
+            if opts.offset is not None:
+                start = opts.offset
+            else:
+                start = state.get('offset:next', 0)
+            outp.printf(f'Starting the dump from offs={start}')
+
+            rotate_size = opts.rotate_size
+            hashes_iter = axon.hashes(start)
+            last_offset = start
+            tar = None
+            tar_path = None
+            file_start = start
+            file_blobcount = 0
+            tar_size = 0
+            for_open = True
+
+            try:
+                async for (offs, (sha256, size)) in hashes_iter:
+                    if for_open:
+                        tar_path = os.path.join(opts.outdir, getTarName(celliden, offs, -1))
+                        tar = tarfile.open(tar_path, 'w:gz')
+                        file_start = offs
+                        tar_size = 0
+                        file_blobcount = 0
+                        for_open = False
+
+                    last_offset = offs
+                    file_blobcount += 1
+                    sha2hex = s_common.ehex(sha256)
+                    outp.printf(f'Dumping blob {sha2hex} (size={size}, offs={offs})')
+
+                    with tempfile.SpooledTemporaryFile(max_size=MAX_SPOOL_SIZE, mode='w+b', dir=opts.outdir) as tmpf:
+                        total = 0
+                        async for byts in axon.get(sha256):
+                            tmpf.write(byts)
+                            total += len(byts)
+                        if total != size:
+                            raise s_exc.BadDataValu(mesg=f'Blob size mismatch for {sha2hex}: expected {size}, got {total}')
+                        tmpf.flush()
+                        tmpf.seek(0)
+                        tarinfo = tarfile.TarInfo(name=f"{sha2hex}.blob")
+                        tarinfo.size = size
+                        tar.addfile(tarinfo, tmpf)
+                        tar_size += size
+
+                        if tar_size >= rotate_size:
+                            outp.printf(f'Rotating to new .tar.gz file at offset {offs + 1}')
+                            tar.close()
+                            final_name = os.path.join(opts.outdir, getTarName(celliden, file_start, offs + 1))
+                            os.rename(tar_path, final_name)
+                            tar = None
+                            tar_path = None
+                            for_open = True
+
+                if tar is not None:
+                    tar.close()
+                    final_name = os.path.join(opts.outdir, getTarName(celliden, file_start, last_offset + 1))
+                    os.rename(tar_path, final_name)
+                    tar = None
+                    tar_path = None
+
+                state['offset:next'] = last_offset + 1
+                s_common.yamlsave(state, statefile_path)
+
+            finally:
+                if tar is not None:
+                    tar.close()
+                    if tar_path and os.path.isfile(tar_path):
+                        os.remove(tar_path)
+
+    except s_exc.SynErr as exc:
+        mesg = exc.get('mesg')
+        return (False, mesg)
+
+    except Exception as e:
+        mesg = f'Error {e} dumping blobs from Axon url: {opts.url}'
+        return (False, mesg)
+
+    return (True, None)
+
+async def main(argv, outp=s_output.stdout):
+    pars = s_cmd.Parser(prog='synapse.tools.axon.dump', outp=outp, description=descr)
+    pars.add_argument('--url', default='cell:///vertex/storage', help='Telepath URL for the Axon.')
+    pars.add_argument('--offset', type=int, default=None, help='Starting offset in the Axon.')
+    pars.add_argument('--rotate-size', type=int, default=DEFAULT_ROTATE_SIZE,
+                      help='Rotate to a new .blobs file after the current file exceeds this size in bytes (default: 4GB). '
+                           'Note: files may exceed this size if a single blob is larger than the remaining space.')
+    pars.add_argument('--statefile', type=str, default=None,
+                      help='Path to the state tracking file for the Axon dump.')
+    pars.add_argument('outdir', help='Directory to dump tar.gz files (required).')
+
+    opts = pars.parse_args(argv)
+
+    async with s_telepath.withTeleEnv():
+        (ok, mesg) = await dumpBlobs(opts, outp)
+        if not ok:
+            outp.printf(f'ERROR: {mesg}')
+            return 1
+
+    outp.printf('Successfully dumped blobs.')
+
+    return 0
+
+if __name__ == '__main__':  # pragma: no cover
+    s_cmd.exitmain(main)

synapse/tools/axon/load.py

@@ -0,0 +1,89 @@
+import tarfile
+
+import synapse.exc as s_exc
+import synapse.common as s_common
+import synapse.telepath as s_telepath
+
+import synapse.lib.cmd as s_cmd
+import synapse.lib.const as s_const
+import synapse.lib.output as s_output
+
+descr = '''
+Load blobs into a Synapse Axon.
+'''
+
+DEFAULT_CHUNK_SIZE = s_const.mebibyte * 16
+
+async def loadBlobs(opts, outp, tarfiles):
+
+    try:
+        async with await s_telepath.openurl(opts.url) as axon:
+
+            cellinfo = await axon.getCellInfo()
+            celltype = cellinfo['cell']['type']
+            if "axon" not in celltype.lower():
+                mesg = f'Axon load tool only works on axons, not {celltype}'
+                raise s_exc.TypeMismatch(mesg=mesg)
+
+            for tarfile_path in tarfiles:
+                outp.printf(f'Processing tar archive: {tarfile_path}')
+                with tarfile.open(tarfile_path, 'r:gz') as tar:
+                    for member in tar:
+                        if not member.name.endswith('.blob'):
+                            continue
+                        sha2hex = member.name[:-5]
+                        try:
+                            sha256 = s_common.uhex(sha2hex)
+                        except Exception:
+                            outp.printf(f"Skipping invalid blob filename: {member.name}")
+                            continue
+                        if await axon.has(sha256):
+                            outp.printf(f"Skipping existing blob {sha2hex}")
+                            continue
+                        outp.printf(f"Loading blob {sha2hex} (size={member.size})")
+                        try:
+                            fobj = tar.extractfile(member)
+                        except OSError as e:
+                            outp.printf(f"WARNING: Error extracting {member.name}: {e}")
+                            continue
+                        if fobj is None:
+                            outp.printf(f"Failed to extract {member.name} from tar archive.")
+                            continue
+                        async with await axon.upload() as upfd:
+                            while True:
+                                chunk = fobj.read(DEFAULT_CHUNK_SIZE)
+                                if not chunk:
+                                    break
+                                await upfd.write(chunk)
+                            await upfd.save()
+
+    except s_exc.SynErr as exc:
+        return (False, s_exc.reprexc(exc))
+
+    except Exception as e:
+        mesg = f'Error {e} loading blobs into Axon url: {opts.url}'
+        return (False, mesg)
+
+    return (True, None)
+
+async def main(argv, outp=s_output.stdout):
+    pars = s_cmd.Parser(prog='synapse.tools.axon.load', outp=outp, description=descr)
+    pars.add_argument('--url', default='cell:///vertex/storage', help='Telepath URL for the Axon.')
+    pars.add_argument('files', nargs='+', help='List of .tar.gz files to import from.')
+
+    opts = pars.parse_args(argv)
+
+    tarfiles = sorted([f for f in opts.files if f.endswith('.tar.gz')])
+
+    async with s_telepath.withTeleEnv():
+        (ok, mesg) = await loadBlobs(opts, outp, tarfiles)
+        if not ok:
+            outp.printf(f'ERROR: {mesg}')
+            return 1
+
+    outp.printf('Successfully loaded blobs.')
+
+    return 0
+
+if __name__ == '__main__':  # pragma: no cover
+    s_cmd.exitmain(main)