diff --git a/kartograf/cli.py b/kartograf/cli.py new file mode 100644 index 0000000..e9d32e2 --- /dev/null +++ b/kartograf/cli.py @@ -0,0 +1,101 @@ +"""CLI functionality for Kartograf.""" + +import argparse +import os +import sys +import kartograf +from kartograf.kartograf import Kartograf + + +def create_parser(): + parser = argparse.ArgumentParser() + + # Command which program Kartograf should run + subparsers = parser.add_subparsers(dest="command") + + parser_map = subparsers.add_parser("map") + + # Write extended logging information to a debug.log file + # TODO: This is always enabled for now but this should be changed when the + # tool is more stable + parser_map.add_argument("-d", "--debug", action="store_true", default=True) + # Delete artifacts from building the requested map + parser_map.add_argument("-c", "--cleanup", action="store_true", default=False) + + parser_map.add_argument("-irr", "--irr", action="store_true", default=False) + parser_map.add_argument("-rv", "--routeviews", action="store_true", default=False) + + # Reproduce a map file from a previous run by using provided input files. + # The input related arguments are ignored if this flag is set. All provided + # input files will be used if they are present in the reproduce folder and + # in the right format. The epoch argument is required for this to work. + parser_map.add_argument("-r", "--reproduce", type=str, default=None) + parser_map.add_argument("-t", "--epoch", type=str, default=None) + + # Waits until the provided epoch is reached before starting the map + parser_map.add_argument("-w", "--wait", type=str, default=None) + + # Ignore ASNs too high to be compatible with the current asmap encoder + # Set to 0 to disable + parser_map.add_argument("-me", "--max_encode", type=int, default=33521664) + + # TODO: + # Save the final output file in a different location that the default out + # folder + # parser_map.add_argument("-o", "--output", action="store_true", default=os.getcwd()) + + # TODO: + # Filter RPKI and IRR data by checking against RIPE RIS and Routeviews data + # and removing all entries that have not been seen announced to those + # services. + # parser_map.add_argument("-f", "--announced_filter", action="store_true", default=False) + + # TODO: + # Include multiple ASNs that validate correctly for the same prefix. + # parser_map.add_argument("-m", "--multi_map", action="store_true", default=False) + + parser_merge = subparsers.add_parser("merge") + + parser_merge.add_argument("-b", "--base", default=f"{os.getcwd()}/base_file.txt") + parser_merge.add_argument("-e", "--extra", default=f"{os.getcwd()}/extra_file.txt") + parser_merge.add_argument("-o", "--output", default=f"{os.getcwd()}/out_file.txt") + + parser_cov = subparsers.add_parser("cov") + + # IP prefix to ASN map to be used for the coverage report. + parser_cov.add_argument("map", type=argparse.FileType("r")) + + # List of IPs to be used for the coverage report. + parser_cov.add_argument("list", type=argparse.FileType("r")) + + parser.add_argument( + "--version", + action="version", + version=f"%(prog)s (version {kartograf.__version__})", + ) + + return parser + + +def main(args=None): + parser = create_parser() + args = parser.parse_args(args) + + if args.command == "run": + if args.reproduce and not args.epoch: + parser.error("--reproduce is required when --epoch is set.") + elif not args.epoch and args.reproduce: + parser.error("--epoch is required when --reproduce is set.") + + if args.wait and args.reproduce: + parser.error("--reproduce is not compatible with --wait.") + + if args.command == "map": + Kartograf.map(args) + elif args.command == "cov": + Kartograf.cov(args) + elif args.command == "merge": + Kartograf.merge(args) + else: + parser.print_help() + sys.exit("Please provide a command.") diff --git a/run b/run index f7c3d1d..40cd424 100755 --- a/run +++ b/run @@ -5,98 +5,7 @@ and IRR as the secondary data source. IRR data is only used if the RPKI data does not cover the respective IP space. """ -import argparse -import os -import sys - -from kartograf.kartograf import Kartograf -import kartograf - +from kartograf.cli import main if __name__ == "__main__": - parser = argparse.ArgumentParser() - - # Command which program Kartograf should run - subparsers = parser.add_subparsers(dest="command") - - parser_map = subparsers.add_parser("map") - - # Write extended logging information to a debug.log file - # TODO: This is always enabled for now but this should be changed when the - # tool is more stable - parser_map.add_argument("-d", "--debug", action="store_true", default=True) - - # Delete artifacts from building the requested map - parser_map.add_argument("-c", "--cleanup", action="store_true", default=False) - - parser_map.add_argument("-irr", "--irr", action="store_true", default=False) - parser_map.add_argument("-rv", "--routeviews", action="store_true", default=False) - - # Reproduce a map file from a previous run by using provided input files. - # The input related arguments are ignored if this flag is set. All provided - # input files will be used if they are present in the reproduce folder and - # in the right format. The epoch argument is required for this to work. - parser_map.add_argument("-r", "--reproduce", type=str, default=None) - parser_map.add_argument("-t", "--epoch", type=str, default=None) - - # Waits until the provided epoch is reached before starting the map - parser_map.add_argument("-w", "--wait", type=str, default=None) - - # Ignore ASNs too high to be compatible with the current asmap encoder - # Set to 0 to disable - parser_map.add_argument("-me", "--max_encode", type=int, default=33521664) - - # TODO: - # Save the final output file in a different location that the default out - # folder - # parser_map.add_argument("-o", "--output", action="store_true", default=os.getcwd()) - - # TODO: - # Filter RPKI and IRR data by checking against RIPE RIS and Routeviews data - # and removing all entries that have not been seen announced to those - # services. - # parser_map.add_argument("-f", "--announced_filter", action="store_true", default=False) - - # TODO: - # Include multiple ASNs that validate correctly for the same prefix. - # parser_map.add_argument("-m", "--multi_map", action="store_true", default=False) - - parser_merge = subparsers.add_parser("merge") - - parser_merge.add_argument("-b", "--base", default=f'{os.getcwd()}/base_file.txt') - parser_merge.add_argument("-e", "--extra", default=f'{os.getcwd()}/extra_file.txt') - parser_merge.add_argument("-o", "--output", default=f'{os.getcwd()}/out_file.txt') - - parser_cov = subparsers.add_parser("cov") - - # IP prefix to ASN map to be used for the coverage report. - parser_cov.add_argument("map", type=argparse.FileType('r')) - - # List of IPs to be used for the coverage report. - parser_cov.add_argument("list", type=argparse.FileType('r')) - - parser.add_argument( - "--version", - action="version", - version="%(prog)s (version {version})".format(version=kartograf.__version__)) - - args = parser.parse_args() - - if args.command == 'run': - if args.reproduce and not args.epoch: - parser.error('--reproduce is required when --epoch is set.') - elif not args.epoch and args.reproduce: - parser.error('--epoch is required when --reproduce is set.') - - if args.wait and args.reproduce: - parser.error('--reproduce is not compatible with --wait.') - - if args.command == "map": - Kartograf.map(args) - elif args.command == "cov": - Kartograf.cov(args) - elif args.command == "merge": - Kartograf.merge(args) - else: - parser.print_help() - sys.exit("Please provide a command.") + main() diff --git a/tests/test_cli.py b/tests/test_cli.py index 5c09bf3..a1d525a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,5 +1,7 @@ import pytest -from run import parser +from kartograf.cli import create_parser + +parser = create_parser() def test_map_command(): args = parser.parse_args(['map']) @@ -45,4 +47,4 @@ def test_version_flag(capsys): parser.parse_args(['--version']) assert excinfo.value.code == 0 captured = capsys.readouterr() - assert "version" in captured.out \ No newline at end of file + assert "version" in captured.out