signal

import argparse
import subprocess
import shlex
import signal


# Function to execute a binary and kill it when a needle is found in its output or after a timeout
# Parameters:
#   binary_path: Path to the binary to execute
#   binary_parameters: Parameters to pass to the binary
#   needle: Needle to match in the output
#   timeout: Timeout after which to kill the process
# Returns:
#   The return code of the binary
# Notes:
#   This function uses subprocess.Popen to execute the binary in a sandbox and capture its output.
#   It then loops through the output lines and checks if the needle is found in the output.
#   If the needle is found, the process is killed and the function returns. If the needle is not found,
#   the function waits for the process to finish and returns its return code.
#   Since popen is used, it does not accept a timeout parameter.
#   Instead, the timeout is implemented by using the timeout command to run the binary.
#   We could not use the timeout parameter of subprocess.run because it blocks the execution of the script
#   until the process finishes.
#   This means that we would not be able to capture the output of the process.
def using_popen(binary_path, binary_parameters, needle, timeout):
    # Define the command to run the binary file
    command = f"timeout {timeout} {binary_path} {binary_parameters}"
    # Use subprocess to execute the command in a sandbox and capture its output
    process = subprocess.Popen(shlex.split(command), stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    # Loop through the output lines and process them
    while True:
        output = process.stderr.readline().decode().strip()
        if output == "" and process.poll() is not None:
            break
        if output != "":
            # Process the output here
            print(output)
            # Check if the output contains the keyword
            if needle in output:
                # Kill the process if the keyword is found
                process.send_signal(signal.SIGINT)
                break
    return process.returncode


if __name__ == "__main__":
    # Parse command line arguments
    parser = argparse.ArgumentParser(
        description='Execute a binary and kill it either when a needle is found in its output or after a timeout')
    parser.add_argument('-e', '--executable', type=str, help='Path to the executable to run')
    parser.add_argument('-p', '--parameters', type=str, help='Parameters to pass to the executable')
    parser.add_argument('-n', '--needle', type=str, help='Needle to match in the output')
    parser.add_argument('-t', '--timeout', type=str, help='Timeout after which to kill the process')
    # Example usage
    # python main.py -e "python" -p "-u -m http.server" -n "404" -t "15s"

    args = parser.parse_args()
    # Execute the binary and capture its output
    return_code = using_popen(args.executable, args.parameters, args.needle, args.timeout)

    # Print the http server's return code
    print("Http server returned with code:", return_code)