diff --git a/code/iottb/__main__.py b/code/iottb/__main__.py index a14754a..1512c70 100644 --- a/code/iottb/__main__.py +++ b/code/iottb/__main__.py @@ -1,9 +1,14 @@ #!/usr/bin/env python3 import argparse +from os import environ +from pathlib import Path -from iottb.subcommands.capture import setup_capture_parser +from iottb.logger import logger from iottb.subcommands.add_device import setup_init_device_root_parser +from iottb.subcommands.capture import setup_capture_parser from iottb.utils.tcpdump_utils import list_interfaces +from definitions import IOTTB_HOME_ABS, ReturnCodes + ###################### # Argparse setup @@ -28,7 +33,37 @@ def setup_argparse(): return root_parser +def check_iottb_env(): + # This makes the option '--root-dir' obsolescent # TODO How to streamline this?\ + try: + iottb_home = environ['IOTTB_HOME'] # TODO WARN implicit declaration of env var name! + except KeyError: + logger.error(f"Environment variable 'IOTTB_HOME' is not set." + f"Setting environment variable 'IOTTB_HOME' to '~/{IOTTB_HOME_ABS}'") + environ['IOTTB_HOME'] = IOTTB_HOME_ABS + finally: + if not Path(IOTTB_HOME_ABS).exists(): + print(f'"{IOTTB_HOME_ABS}" does not exist.') + response = input('Do you want to create it now? [y/N]') + logger.debug(f'response: {response}') + if response.lower() != 'y': + logger.debug(f'Not creating "{environ['IOTTB_HOME']}"') + print('TODO') + print("Aborting execution...") + return ReturnCodes.ABORTED + else: + print(f'Creating "{environ['IOTTB_HOME']}"') + Path(IOTTB_HOME_ABS).mkdir(parents=True, + exist_ok=False) # Should always work since in 'not exist' code path + return ReturnCodes.OK + logger.info(f'"{IOTTB_HOME_ABS}" exists.') + # TODO: Check that it is a valid iottb dir or can we say it is valid by definition if? + return ReturnCodes.OK + + def main(): + if check_iottb_env() != ReturnCodes.OK: + exit(ReturnCodes.ABORTED) parser = setup_argparse() args = parser.parse_args() print(args) diff --git a/code/iottb/definitions.py b/code/iottb/definitions.py index 3e6e6a9..c962e76 100644 --- a/code/iottb/definitions.py +++ b/code/iottb/definitions.py @@ -1,6 +1,17 @@ from datetime import datetime from enum import Flag, unique, global_enum +from pathlib import Path + +''' +Defining IOTTB_HOME_ABS here implies that it be immutable. +It is used here so that one could configure it. +But after its used in __man__ this cannot be relied upon. +''' +IOTTB_HOME_ABS = Path().home() / 'IOTTB.db' + +# TODO maybe wrap this into class to make it easier to pass along to different objects +# But will need more refactoring DEVICE_METADATA_FILE = 'device_metadata.json' CAPTURE_METADATA_FILE = 'capture_metadata.json' TODAY_DATE_STRING = datetime.now().strftime('%d%b%Y').lower() # TODO convert to function in utils or so @@ -24,3 +35,7 @@ class ReturnCodes(Flag): FILE_ALREADY_EXISTS = 5 INVALID_ARGUMENT = 6 INVALID_ARGUMENT_VALUE = 7 + + +def iottb_home_abs(): + return None