83 lines
3.0 KiB
Python
83 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
from os import environ
|
|
from pathlib import Path
|
|
|
|
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
|
|
######################
|
|
def setup_argparse():
|
|
# create top level parser
|
|
root_parser = argparse.ArgumentParser(prog='iottb')
|
|
subparsers = root_parser.add_subparsers(title='subcommands', required=True, dest='command')
|
|
|
|
# shared options
|
|
root_parser.add_argument('--verbose', '-v', action='count', default=0)
|
|
# configure subcommands
|
|
setup_capture_parser(subparsers)
|
|
setup_init_device_root_parser(subparsers)
|
|
|
|
# Utility to list interfaces directly with iottb instead of relying on external tooling
|
|
|
|
interfaces_parser = subparsers.add_parser('list-interfaces', aliases=['li', 'if'],
|
|
help='List available network interfaces.')
|
|
interfaces_parser.set_defaults(func=list_interfaces)
|
|
|
|
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)
|
|
if args.command:
|
|
try:
|
|
args.func(args)
|
|
except KeyboardInterrupt:
|
|
print('Received keyboard interrupt. Exiting...')
|
|
exit(1)
|
|
except Exception as e:
|
|
print(f'Error: {e}')
|
|
# create_capture_directory(args.device_name)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|