4a6fd338e0
* replace 'timeout' helper with async python script; allow hub preload in func tests; improve item removal * func tests: cscli hub update/upgrade * docker test update * Update docker entrypoint to disable items with --force The --force flag was not transmitted to cscli, but is required after the hub refact to disable items inside installed collections
117 lines
3.9 KiB
Python
Executable file
117 lines
3.9 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
|
|
import asyncio
|
|
import argparse
|
|
import os
|
|
import re
|
|
import signal
|
|
import sys
|
|
|
|
DEFAULT_TIMEOUT = 30
|
|
|
|
# TODO: signal handler to terminate spawned process group when wait-for is killed
|
|
# TODO: better return codes esp. when matches are found
|
|
# TODO: multiple patterns (multiple out, err, both)
|
|
# TODO: print unmatched patterns
|
|
|
|
|
|
async def terminate(p):
|
|
# Terminate the process group (shell, crowdsec plugins)
|
|
try:
|
|
os.killpg(os.getpgid(p.pid), signal.SIGTERM)
|
|
except ProcessLookupError:
|
|
pass
|
|
|
|
|
|
async def monitor(cmd, args, want_out, want_err, timeout):
|
|
"""Monitor a process and terminate it if a pattern is matched in stdout or stderr.
|
|
|
|
Args:
|
|
cmd: The command to run.
|
|
args: A list of arguments to pass to the command.
|
|
stdout: A regular expression pattern to search for in stdout.
|
|
stderr: A regular expression pattern to search for in stderr.
|
|
timeout: The maximum number of seconds to wait for the process to terminate.
|
|
|
|
Returns:
|
|
The exit code of the process.
|
|
"""
|
|
|
|
status = None
|
|
|
|
async def read_stream(p, stream, outstream, pattern):
|
|
nonlocal status
|
|
if stream is None:
|
|
return
|
|
while True:
|
|
line = await stream.readline()
|
|
if line:
|
|
line = line.decode('utf-8')
|
|
outstream.write(line)
|
|
if pattern and pattern.search(line):
|
|
await terminate(process)
|
|
# this is nasty.
|
|
# if we timeout, we want to return a different exit code
|
|
# in case of a match, so that the caller can tell
|
|
# if the application was still running.
|
|
# XXX: still not good for match found, but return code != 0
|
|
if timeout != DEFAULT_TIMEOUT:
|
|
status = 128
|
|
else:
|
|
status = 0
|
|
break
|
|
else:
|
|
break
|
|
|
|
process = await asyncio.create_subprocess_exec(
|
|
cmd,
|
|
*args,
|
|
# capture stdout
|
|
stdout=asyncio.subprocess.PIPE,
|
|
# capture stderr
|
|
stderr=asyncio.subprocess.PIPE,
|
|
# disable buffering
|
|
bufsize=0,
|
|
# create a new process group
|
|
# (required to kill child processes when cmd is a shell)
|
|
preexec_fn=os.setsid)
|
|
|
|
out_regex = re.compile(want_out) if want_out else None
|
|
err_regex = re.compile(want_err) if want_err else None
|
|
|
|
# Apply a timeout
|
|
try:
|
|
await asyncio.wait_for(
|
|
asyncio.wait([
|
|
asyncio.create_task(process.wait()),
|
|
asyncio.create_task(read_stream(process, process.stdout, sys.stdout, out_regex)),
|
|
asyncio.create_task(read_stream(process, process.stderr, sys.stderr, err_regex))
|
|
]), timeout)
|
|
if status is None:
|
|
status = process.returncode
|
|
except asyncio.TimeoutError:
|
|
await terminate(process)
|
|
status = 241
|
|
|
|
# Return the same exit code, stdout and stderr as the spawned process
|
|
return status
|
|
|
|
|
|
async def main():
|
|
parser = argparse.ArgumentParser(
|
|
description='Monitor a process and terminate it if a pattern is matched in stdout or stderr.')
|
|
parser.add_argument('cmd', help='The command to run.')
|
|
parser.add_argument('args', nargs=argparse.REMAINDER, help='A list of arguments to pass to the command.')
|
|
parser.add_argument('--out', default='', help='A regular expression pattern to search for in stdout.')
|
|
parser.add_argument('--err', default='', help='A regular expression pattern to search for in stderr.')
|
|
parser.add_argument('--timeout', type=float, default=DEFAULT_TIMEOUT)
|
|
args = parser.parse_args()
|
|
|
|
exit_code = await monitor(args.cmd, args.args, args.out, args.err, args.timeout)
|
|
|
|
sys.exit(exit_code)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
asyncio.run(main())
|