Refactor and add Logger

code/iottb/subcommands/add_device.py

@@ -0,0 +1,59 @@
+import pathlib
+
+from iottb import definitions
+from iottb.definitions import DEVICE_METADATA_FILE, ReturnCodes
+from iottb.models.device_metadata_model import DeviceMetadata
+from iottb.utils.device_metadata_utils import *
+
+
+def setup_init_device_root_parser(subparsers):
+    parser = subparsers.add_parser("add-device", aliases=["add-device-root", "add"])
+    parser.add_argument("--root_dir", type=pathlib.Path, default=pathlib.Path.cwd())
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument("--guided", action="store_true", help="Guided setup", default=False)
+    group.add_argument("--name", action="store", type=str, help="name of device")
+    parser.set_defaults(func=handle_add)
+
+
+def handle_add(args):
+    print("Entered add-device-root")
+
+    if args.guided:
+        metadata = guided_setup(args.root_dir)
+    else:
+        device_name = args.name
+        args.root_dir.mkdir(parents=True, exist_ok=True)
+        args.root_dir.chdir()
+        metadata = DeviceMetadata(device_name, args.root_dir)
+
+    file_path = args.root_dir / DEVICE_METADATA_FILE
+    response = input(f"Confirm device metadata: {metadata.model_dump()} [y/N]")
+    if response.lower() not in definitions.AFFIRMATIVE_USER_RESPONSE.add(""):
+        configure_metadata()
+        assert False, "TODO implement dynamic setup"
+    assert metadata.model_dump() != ""
+    if metadata.save_to_json(file_path) == ReturnCodes.FILE_ALREADY_EXISTS:
+        print("Directory already contains a device metadata file. Aborting operation.")
+        return ReturnCodes.ABORTED
+    assert Path(file_path).exists(), f"{file_path} does not exist"
+    return ReturnCodes.SUCCESS
+
+
+def configure_metadata():
+    pass
+
+
+def guided_setup(device_root) -> DeviceMetadata:
+    response = "N"
+    device_name = ""
+    while response.upper() == "N":
+        device_name = input("Please enter name of device: ")
+        if device_name == "" or device_name is None:
+            print("Name cannot be empty")
+        response = input(f"Confirm device name: {device_name} [y/N] ")
+
+    assert response.lower() in definitions.AFFIRMATIVE_USER_RESPONSE.add(""), f"{response.upper()} not supported"
+    assert device_name != ""
+    assert device_name is not None
+    return DeviceMetadata(device_name, device_root)
+

code/iottb/subcommands/capture.py

@@ -0,0 +1,170 @@
+import subprocess
+from pathlib import Path
+
+from iottb.definitions import *
+from iottb.models.capture_metadata_model import CaptureMetadata
+from iottb.models.device_metadata_model import DeviceMetadata, dir_contains_device_metadata
+from iottb.utils.capture_utils import get_capture_src_folder, make_capture_src_folder
+
+
+def setup_capture_parser(subparsers):
+    parser = subparsers.add_parser('sniff', help='Sniff packets with tcpdump')
+    # metadata args
+    parser.add_argument("-a", "--ip-address", help="IP address of the device to sniff", dest="device_ip")
+    # tcpdump args
+    parser.add_argument("device_root", help="Root folder for device to sniff",
+                        type=Path, default=Path.cwd())
+    parser.add_argument("-s", "--safe", help="Ensure correct device root folder before sniffing", action="store_true")
+    parser.add_argument("--app", help="Application name to sniff", dest="app_name", default=None)
+
+    parser_sniff_tcpdump = parser.add_argument_group('tcpdump arguments')
+    parser_sniff_tcpdump.add_argument("-i", "--interface", help="Interface to capture on.", dest="capture_interface",
+                                      required=True)
+    parser_sniff_tcpdump.add_argument("-I", "--monitor-mode", help="Put interface into monitor mode",
+                                      action="store_true")
+    parser_sniff_tcpdump.add_argument("-n", help="Deactivate name resolution. True by default.",
+                                      action="store_true", dest="no_name_resolution")
+    parser_sniff_tcpdump.add_argument("-#", "--number",
+                                      help="Print packet number at beginning of line. True by default.",
+                                      action="store_true")
+    parser_sniff_tcpdump.add_argument("-e", help="Print link layer headers. True by default.",
+                                      action="store_true", dest="print_link_layer")
+    parser_sniff_tcpdump.add_argument("-t", action="count", default=0,
+                                      help="Please see tcpdump manual for details. Unused by default.")
+
+    cap_size_group = parser.add_mutually_exclusive_group(required=False)
+    cap_size_group.add_argument("-c", "--count", type=int, help="Number of packets to capture.", default=1000)
+    cap_size_group.add_argument("--mins", type=int, help="Time in minutes to capture.", default=1)
+
+    parser.set_defaults(func=handle_capture)
+
+
+def cwd_is_device_root_dir() -> bool:
+    device_metadata_file = Path.cwd() / DEVICE_METADATA_FILE
+    return device_metadata_file.is_file()
+
+
+def start_guided_device_root_dir_setup():
+    assert False, "Not implemented"
+
+
+def handle_metadata():
+    assert not cwd_is_device_root_dir()
+    print(f"Unable to find {DEVICE_METADATA_FILE} in current working directory")
+    print("You need to setup a device root directory before using this command")
+    response = input("Would you like to be guided through the setup? [y/n]")
+    if response.lower() == "y":
+        start_guided_device_root_dir_setup()
+    else:
+        print("'iottb init-device-root --help' for more information.")
+        exit(ReturnCodes.ABORTED)
+    # device_id = handle_capture_metadata()
+    return ReturnCodes.SUCCESS
+
+
+def get_device_metadata_from_file(device_metadata_filename: Path) -> str:
+    assert device_metadata_filename.is_file(), f"Device metadata file '{device_metadata_filename} does not exist"
+    device_metadata = DeviceMetadata.load_from_json(device_metadata_filename)
+    return device_metadata
+
+
+def run_tcpdump(cmd):
+    # TODO: Maybe specify files for stout and stderr
+    try:
+        p = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        if p.returncode != 0:
+            print(f"Error running tcpdump {p.stderr}")
+        else:
+            print(f"tcpdump run successfully\n: {p.stdout}")
+    except KeyboardInterrupt:
+        pass
+
+
+def handle_capture(args):
+    assert args.device_root is not None, f"Device root directory is required"
+    assert dir_contains_device_metadata(args.device_root), f"Device metadata file '{args.device_root}' does not exist"
+    # get device metadata
+    if args.safe and not dir_contains_device_metadata(args.device_root):
+        print(f"Supplied folder contains no device metadata. "
+              f"Please setup a device root directory before using this command")
+        exit(ReturnCodes.ABORTED)
+    elif dir_contains_device_metadata(args.device_root):
+        device_metadata_filename = args.device_root / DEVICE_METADATA_FILE
+        device_data = DeviceMetadata.load_from_json(device_metadata_filename)
+    else:
+        name = input("Please enter a device name: ")
+        args.device_root.mkdir(parents=True, exist_ok=True)
+        device_data = DeviceMetadata(name, args.device_root)
+    # start constructing environment for capture
+    capture_dir = get_capture_src_folder(args.device_root)
+    make_capture_src_folder(capture_dir)
+    capture_metadata = CaptureMetadata(device_data, capture_dir)
+
+    capture_metadata.set_interface(args.capture_interface)
+    cmd = ['sudo', 'tcpdump', '-i', args.capture_interface]
+    cmd = build_tcpdump_args(args, cmd, capture_metadata)
+    capture_metadata.set_tcpdump_command(cmd)
+
+    print('Executing: ' + ' '.join(cmd))
+
+    # run capture
+    try:
+        start_time = datetime.now().strftime('%H:%M:%S')
+        run_tcpdump(cmd)
+        stop_time = datetime.now().strftime('%H:%M:%S')
+        capture_metadata.set_start_time(start_time)
+        capture_metadata.set_stop_time(stop_time)
+    except KeyboardInterrupt:
+        print("Received keyboard interrupt.")
+        exit(ReturnCodes.ABORTED)
+    except subprocess.CalledProcessError as e:
+        print(f"Failed to capture packet: {e}")
+        exit(ReturnCodes.FAILURE)
+    except Exception as e:
+        print(f"Failed to capture packet: {e}")
+        exit(ReturnCodes.FAILURE)
+
+    return ReturnCodes.SUCCESS
+
+
+def build_tcpdump_args(args, cmd, capture_metadata: CaptureMetadata):
+    if args.monitor_mode:
+        cmd.append('-I')
+    if args.no_name_resolution:
+        cmd.append('-n')
+    if args.number:
+        cmd.append('-#')
+    if args.print_link_layer:
+        cmd.append('-e')
+
+    if args.count:
+        cmd.append('-c')
+        cmd.append(str(args.count))
+    elif args.mins:
+        assert False, "Unimplemented option"
+
+    if args.app_name is not None:
+        capture_metadata.set_app_name(args.app_name)
+
+    capture_metadata.build_capture_file_name()
+    cmd.append('-w')
+    cmd.append(capture_metadata.get_capfile_name())
+
+    if args.safe:
+        cmd.append(f'host {args.device_ip}')  # if not specified, filter 'any' implied by tcpdump
+        capture_metadata.set_device_ip_address(args.device_ip)
+
+    return cmd
+
+
+# def capture_file_cmd(args, cmd, capture_dir, capture_metadata: CaptureMetadata):
+#     capture_file_prefix = capture_metadata.get_device_metadata().get_device_short_name()
+#     if args.app_name is not None:
+#         capture_file_prefix = args.app_name
+#         capture_metadata.set_app(args.app_name)
+#     capfile_name = capture_file_prefix + "_" + str(capture_metadata.get_capture_id()) + ".pcap"
+#     capture_metadata.set_capture_file(capfile_name)
+#     capfile_abs_path = capture_dir / capfile_name
+#     capture_metadata.set_capture_file(capfile_name)
+#     cmd.append('-w')
+#     cmd.append(str(capfile_abs_path))