github.com/crowdsecurity/crowdsec@v1.6.1/test/bin/wait-for (about)

     1  #!/usr/bin/env python3
     2  
     3  import asyncio
     4  import argparse
     5  import os
     6  import re
     7  import signal
     8  import sys
     9  
    10  DEFAULT_TIMEOUT = 30
    11  
    12  # TODO: signal handler to terminate spawned process group when wait-for is killed
    13  # TODO: better return codes esp. when matches are found
    14  # TODO: multiple patterns (multiple out, err, both)
    15  # TODO: print unmatched patterns
    16  
    17  
    18  async def terminate(p):
    19      # Terminate the process group (shell, crowdsec plugins)
    20      try:
    21          os.killpg(os.getpgid(p.pid), signal.SIGTERM)
    22      except ProcessLookupError:
    23          pass
    24  
    25  
    26  async def monitor(cmd, args, want_out, want_err, timeout):
    27      """Monitor a process and terminate it if a pattern is matched in stdout or stderr.
    28  
    29      Args:
    30          cmd: The command to run.
    31          args: A list of arguments to pass to the command.
    32          stdout: A regular expression pattern to search for in stdout.
    33          stderr: A regular expression pattern to search for in stderr.
    34          timeout: The maximum number of seconds to wait for the process to terminate.
    35  
    36      Returns:
    37          The exit code of the process.
    38      """
    39  
    40      status = None
    41  
    42      async def read_stream(p, stream, outstream, pattern):
    43          nonlocal status
    44          if stream is None:
    45              return
    46          while True:
    47              line = await stream.readline()
    48              if line:
    49                  line = line.decode('utf-8')
    50                  outstream.write(line)
    51                  if pattern and pattern.search(line):
    52                      await terminate(process)
    53                      # this is nasty.
    54                      # if we timeout, we want to return a different exit code
    55                      # in case of a match, so that the caller can tell
    56                      # if the application was still running.
    57                      # XXX: still not good for match found, but return code != 0
    58                      if timeout != DEFAULT_TIMEOUT:
    59                          status = 128
    60                      else:
    61                          status = 0
    62                      break
    63              else:
    64                  break
    65  
    66      process = await asyncio.create_subprocess_exec(
    67          cmd,
    68          *args,
    69          # capture stdout
    70          stdout=asyncio.subprocess.PIPE,
    71          # capture stderr
    72          stderr=asyncio.subprocess.PIPE,
    73          # disable buffering
    74          bufsize=0,
    75          # create a new process group
    76          # (required to kill child processes when cmd is a shell)
    77          preexec_fn=os.setsid)
    78  
    79      out_regex = re.compile(want_out) if want_out else None
    80      err_regex = re.compile(want_err) if want_err else None
    81  
    82      # Apply a timeout
    83      try:
    84          await asyncio.wait_for(
    85              asyncio.wait([
    86                  asyncio.create_task(process.wait()),
    87                  asyncio.create_task(read_stream(process, process.stdout, sys.stdout, out_regex)),
    88                  asyncio.create_task(read_stream(process, process.stderr, sys.stderr, err_regex))
    89              ]), timeout)
    90          if status is None:
    91              status = process.returncode
    92      except asyncio.TimeoutError:
    93          await terminate(process)
    94          status = 241
    95  
    96      # Return the same exit code, stdout and stderr as the spawned process
    97      return status
    98  
    99  
   100  async def main():
   101      parser = argparse.ArgumentParser(
   102          description='Monitor a process and terminate it if a pattern is matched in stdout or stderr.')
   103      parser.add_argument('cmd', help='The command to run.')
   104      parser.add_argument('args', nargs=argparse.REMAINDER, help='A list of arguments to pass to the command.')
   105      parser.add_argument('--out', default='', help='A regular expression pattern to search for in stdout.')
   106      parser.add_argument('--err', default='', help='A regular expression pattern to search for in stderr.')
   107      parser.add_argument('--timeout', type=float, default=DEFAULT_TIMEOUT)
   108      args = parser.parse_args()
   109  
   110      exit_code = await monitor(args.cmd, args.args, args.out, args.err, args.timeout)
   111  
   112      sys.exit(exit_code)
   113  
   114  
   115  if __name__ == '__main__':
   116      asyncio.run(main())