diff --git a/json_dump/json_dump.py b/json_dump/json_dump.py index ede48221..2254906d 100755 --- a/json_dump/json_dump.py +++ b/json_dump/json_dump.py @@ -1,91 +1,156 @@ #!/usr/bin/python3 -import sys -import os.path -import socket +import argparse import io -import pytrap import json -import optparse - -from optparse import OptionParser -parser = OptionParser(add_help_option=True) -parser.add_option("-i", "--ifcspec", dest="ifcspec", - help="See https://nemea.liberouter.org/trap-ifcspec/", metavar="IFCSPEC") -parser.add_option("-w", dest="filename", - help="Write dump to FILE instead of stdout (overwrite file)", metavar="FILE") -parser.add_option("-a", dest="filename_append", - help="Write dump to FILE instead of stdout (append to file)", metavar="FILE") -parser.add_option("-s", dest="networktarget", - help="Stream data over TCP", metavar="HOST:PORT") -parser.add_option("-I", "--indent", metavar="N", type=int, - help="Pretty-print JSON with indentation set to N spaces. Note that such format can't be read by json_replay module.") -parser.add_option("-v", "--verbose", action="store_true", - help="Set verbose mode - print messages.") -parser.add_option("--noflush", action="store_true", - help="Disable automatic flush of output buffer after writing a record (may improve performance).") - - -# Parse remaining command-line arguments -(options, args) = parser.parse_args() - -# Initialize module -trap = pytrap.TrapCtx() -trap.init(["-i", options.ifcspec]) - -# Open output file -if options.filename and options.filename_append: - sys.stderr.write("Error: -w and -a are mutually exclusive.") - sys.exit(1) -if options.filename: - file = io.FileIO(options.filename, "w") -elif options.filename_append: - file = io.FileIO(options.filename_append, "a") -elif options.networktarget: - addr = options.networktarget.split(":") - if len(addr) != 2: - raise AttributeError("Malformed argument of -s host:port") - s = socket.create_connection((addr[0], int(addr[1]))) - file = socket.SocketIO(s, "w") -else: - file = sys.stdout - -# Set JSON as required data type on input -trap.setRequiredFmt(0, pytrap.FMT_JSON, "") - -stop = False -# Main loop (trap.stop is set to True when SIGINT or SIGTERM is received) -while not stop: - # Read data from input interface - try: - data = trap.recv() - except pytrap.FormatMismatch: - sys.stderr.write("Error: output and input interfaces data type or format mismatch\n") - break - except pytrap.FormatChanged as e: - if options.verbose: - print(trap.getDataFmt(0)) - data = e.data - del(e) - pass - except (pytrap.Terminated, KeyboardInterrupt): - break - - # Check for "end-of-stream" record - if len(data) <= 1: - if options.verbose: - print('Received "end-of-stream" message, going to quit.') - break +import socket +import sys +import time +from typing import Optional, TextIO + +import pytrap + + +def get_parser() -> argparse.ArgumentParser: + """Prepare the argument parser. + + Returns: + An instance of ArgumentParser with ready arguments. + """ + parser = argparse.ArgumentParser(description="Print received JSON messages to stdout or a file.") + parser.add_argument("-i", "--ifcspec", dest="ifcspec", metavar="IFCSPEC", + required=True, + help="See https://nemea.liberouter.org/trap-ifcspec/") + parser.add_argument("-I", "--indent", metavar="N", type=int, + help="Pretty-print JSON with indentation set to N spaces. " + "Note that such format can't be read by json_replay module.") + parser.add_argument("-v", "--verbose", action="store_true", + help="Set verbose mode (print messages).") + parser.add_argument("--noflush", action="store_true", + help="Disable automatic flush of output buffer after writing a " + "record (may improve performance).") + + group = parser.add_mutually_exclusive_group() + group.add_argument("-w", dest="filename", metavar="FILE", + help="Write data to FILE instead of stdout (overwrite file)") + group.add_argument("-a", dest="filename_append", metavar="FILE", + help="Write data to FILE instead of stdout (append to file)") + group.add_argument("-s", dest="networktarget", metavar="HOST:PORT", + help="Send data using a TCP network stream to HOST:PORT") + return parser + + +def connect_socket(address: str, port: int, wait_interval: int = 5) -> TextIO: + """Create a connection to a socket given an address and a port. + + The connection is tried repeatedly until successful. If the connection to the socket + fails, wait 5 seconds and retry. + + Args: + address (str): The address of the destination socket. + port (int): The port of the destination socket. + wait_interval (int): Number of seconds to wait before retrying connection if it fails (default: 5 sec) + + Returns: + TextIO object providing access to opened socket. + """ + last_error = None + while True: + if last_error is None: + print(f"{time.strftime('%F-%T')} Connecting to {address}:{port} ...", file=sys.stderr) + try: + s = socket.create_connection((address, port)) + print(f"{time.strftime('%F-%T')} Connection established.", file=sys.stderr) + last_error = None + return s.makefile("w", encoding="utf-8") + except OSError as e: + # sleep for a while and then reconnect (print error only once or when the error message changes) + if last_error is None or str(e) != last_error: + print(f"{time.strftime('%F-%T')} Connection failed ({e}), retrying every {wait_interval} seconds ...", file=sys.stderr) + last_error = str(e) + time.sleep(wait_interval) + + +def main(): + parser = get_parser() + parsed_args = parser.parse_args() + + address: Optional[str] = None + port: Optional[int] = None + + # Parsing the arguments + if parsed_args.filename: + file = open(parsed_args.filename, "w", encoding="utf-8") + elif parsed_args.filename_append: + file = open(parsed_args.filename_append, "a", encoding="utf-8") + elif parsed_args.networktarget: + try: + address, port = parsed_args.networktarget.split(":") + port = int(port) + except (TypeError, ValueError): + print("Error: malformed argument of -s host:port", file=sys.stderr) + sys.exit(1) + file = connect_socket(address, port) + else: + file = sys.stdout + + # Initialize the PyTrap module + trap = pytrap.TrapCtx() + trap.init(["-i", parsed_args.ifcspec]) + + # Set JSON as required data type on input + trap.setRequiredFmt(0, pytrap.FMT_JSON, "") + + # Main loop + while True: + # Read data from input interface + try: + data = trap.recv() + except pytrap.FormatMismatch: + print("Error: output and input interfaces data type or format mismatch", file=sys.stderr) + break + except pytrap.FormatChanged as e: + if parsed_args.verbose: + print(trap.getDataFmt(0)) + data = e.data + except (pytrap.Terminated, KeyboardInterrupt): + break + + # Check for "end-of-stream" record + if len(data) <= 1: + if parsed_args.verbose: + print('Received "end-of-stream" message, going to quit.') + break - try: # Decode data (and check it's valid JSON) - rec = json.loads(data.decode("utf-8")) - if options.verbose: - print("Message: {0}".format(rec)) - # Print it to file or stdout - file.write(bytes(json.dumps(rec, indent=options.indent) + '\n', "utf-8")) - if not options.noflush: - file.flush() - except ValueError as e: - sys.stderr.write(str(e) + '\n') + try: + rec = json.loads(data.decode("utf-8")) + except ValueError as e: + print(f"ERROR: Received invalid JSON (message skipped): {e}", file=sys.stderr) + continue + + if parsed_args.verbose: + print(f"Message: {format(rec)}") + # Print it to file, stdout, or send to socket + try: + file.write(json.dumps(rec, indent=parsed_args.indent) + '\n') + if not parsed_args.noflush: + file.flush() + except IOError as e: + if parsed_args.networktarget: + print(f"{time.strftime('%F-%T')} Connection error: {e}", file=sys.stderr) + # connection error, try to reconnect and send the message again + file = connect_socket(address, port) + file.write(json.dumps(rec, indent=parsed_args.indent) + '\n') + if not parsed_args.noflush: + file.flush() + else: + raise + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + pass # quietly exit program without traceback