Refactor subcommands, config etc.

This commit is contained in:
Sebastian Lenzlinger
2024-05-05 17:54:21 +02:00
parent 7e3e101987
commit ec08bf71ab
8 changed files with 135 additions and 106 deletions

View File

@@ -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", 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")
metadata = None
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)
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)

View File

@@ -1,8 +1,10 @@
import subprocess
import uuid
from pathlib import Path
from iottb.definitions import *
from iottb.models.device_metadata_model import DeviceMetadata
from iottb.models.capture_metadata_model import CaptureMetadata
def setup_capture_parser(subparsers):
@@ -37,7 +39,7 @@ def setup_capture_parser(subparsers):
def cwd_is_device_root_dir() -> bool:
device_metadata_file = Path.cwd() / DEVICE_METADATA_FILE
return device_metadata_file.exists()
return device_metadata_file.is_file()
def get_user_capture_config():
@@ -61,57 +63,63 @@ def handle_metadata():
return ReturnCodes.SUCCESS
def handle_capture_metadata():
device_metadata_json = Path.cwd() / DEVICE_METADATA_FILE
device_metadata = DeviceMetadata.load_from_json(device_metadata_json)
def get_device_id_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)
device_id = device_metadata.device_id
return device_id
def handle_date_dir():
def setup_capture_folder():
pass
def run_tcpdum(cmd):
subprocess.run(cmd)
try:
p = subprocess.run(cmd)
except KeyboardInterrupt:
pass
def generate_capfile_name():
name = datetime.now().strftime("%Y%m%d_%H%M%S") + uuid.uuid4().hex
def handle_capture(args):
if cwd_is_device_root_dir():
handle_date_dir()
cmd = ['sudo tcpdump', '-i', args.capture_interface]
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:
pass
print('Executing: ' + ' '.join(cmd))
# TODO maybe dump this into file -> put into device metadata
try:
start_time = datetime.now().strftime('%H:%M:%S')
run_tcpdum(cmd)
stop_time = datetime.now().strftime('%H:%M:%S')
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)
assert args.device_root is not None, f"Device root directory is required"
device_metadata_file = Path.cwd() / DEVICE_METADATA_FILE
device_id = get_device_id_from_file(device_metadata_file)
assert device_metadata_file.is_file(), f"Device metadata file '{device_metadata_file} does not exist"
capture_dir = setup_capture_folder(args.dev_root)
cmd = ['sudo', 'tcpdump', '-i', args.capture_interface]
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:
pass
print('Executing: ' + ' '.join(cmd))
# TODO maybe dump this into file -> put into device metadata
try:
start_time = datetime.now().strftime('%H:%M:%S')
run_tcpdum(cmd)
stop_time = datetime.now().strftime('%H:%M:%S')
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
else:
handle_metadata()
return ReturnCodes.SUCCESS

View File

@@ -1,40 +0,0 @@
import pathlib
from iottb.definitions import DEVICE_METADATA_FILE
from iottb.models.device_metadata_model import DeviceMetadata
def setup_init_iottb_root_parser(subparsers):
parser = subparsers.add_parser("init", aliases=["idr"])
parser.add_argument("--root_dir", type=pathlib.Path, default=pathlib.Path.cwd())
group = parser.add_mutually_exclusive_group()
group.add_argument("--dynamic", action="store_true", help="enable guided setup", default=False)
group.add_argument("-n", "--name", action="store", type=str, help="name of device")
parser.set_defaults(func=handle_idr)
def handle_idr(args):
print("Entered iottb initialize-device-root")
root_dir = args.root_dir
device_name = None
if args.dynamic:
response = "N"
while response == "N":
name = input("Please enter name of device: ")
# TODO extended config for other fields like apps, firmware etc.
response = input(f"Confirm device name: {name} [y/N]")
device_name = name
else:
device_name = args.name
root_dir.mkdir(parents=True, exist_ok=True)
root_dir.chdir()
dev_metadata_model = DeviceMetadata(device_name)
file_path = root_dir / device_name / DEVICE_METADATA_FILE
assert not file_path.exists(), f"{file_path} already exists"
if args.dynamic:
response = input(f"Confirm device metadata: {dev_metadata_model.model_dump()} [y/N]")
if response.lower() != "y":
assert False, "TODO implement dynamic setup"
code = dev_metadata_model.save_to_json(file_path)
print(f"Device metadata saved to {file_path}")
return code