diff --git a/bwscanner/configutil.py b/bwscanner/configutil.py index aef3500..60f2c02 100644 --- a/bwscanner/configutil.py +++ b/bwscanner/configutil.py @@ -6,15 +6,28 @@ def read_config(cfg_path): - log.debug('reading config %s' % cfg_path) + log.debug('Reading config %s' % cfg_path) if not config_exists(cfg_path): copy_config(cfg_path) parser = SafeConfigParser() parser.read([cfg_path]) - # FIXME: handle section names - section = 'default' - return dict(parser.items(section)) - + cfg_dict = dict(parser.items('default')) + int_keys = cfg_dict['int_keys'].split(' ') + bool_keys = cfg_dict['bool_keys'].split(' ') + for k in int_keys: + cfg_dict[k] = int(cfg_dict[k]) + for i in bool_keys: + cfg_dict[k] = bool(cfg_dict[k]) + bw_files = dict(parser.items('bw_files')) + cfg_bw_files = {} + for k, v in bw_files.items(): + print(k, v) + if 'm' in k: + number = k.rstrip('m') + size = 1024 * int(number) + cfg_bw_files[size] = (k.upper(), v) + cfg_dict['bw_files'] = cfg_bw_files + return cfg_dict def config_exists(cfg_path): return os.path.isfile(cfg_path) diff --git a/bwscanner/scanner.py b/bwscanner/scanner.py index 39d16dd..fe1f4ec 100755 --- a/bwscanner/scanner.py +++ b/bwscanner/scanner.py @@ -18,19 +18,18 @@ CONFIG_FILE = 'config.ini' LOG_FILE = 'bwscanner.log' -CTX = dict( - default_map=read_config(os.path.join(DATA_DIR, CONFIG_FILE)) -) - class ScanInstance(object): """ Store the configuration and state for the CLI tool. """ - def __init__(self, data_dir): + def __init__(self, data_dir, measurement_dir=None): self.data_dir = data_dir - self.measurement_dir = os.path.join(data_dir, 'measurements') - self.tor_dir = os.path.join(data_dir, 'tor_data') + if measurement_dir is None: + self.measurement_dir = os.path.join(data_dir, 'measurements') + else: + self.measurement_dir = measurement_dir + self.tor_state = None def __repr__(self): return '' % self.data_dir @@ -40,26 +39,19 @@ def __repr__(self): # FIXME: change all options to take defaults from CTX, ie config file? -@click.group(context_settings=CTX) +@click.group() @click.option('--data-dir', type=click.Path(), - default=os.environ.get("BWSCANNER_DATADIR", - CTX.get('data_dir', - click.get_app_dir(APP_NAME))), help='Directory where bwscan should stores its measurements and ' 'other data.') @click.option('-l', '--loglevel', help='The logging level the scanner will use (default: info)', - default=CTX.get('loglevel', 'info'), - type=click.Choice(['debug', 'info', 'warn', 'error', 'critical'])) + type=click.Choice( + ['debug', 'info', 'warn', 'error', 'critical'])) @click.option('-f', '--logfile', type=click.Path(), - help='The file the log will be written to', - default=os.environ.get("BWSCANNER_LOGFILE", - CTX.get('logfile', LOG_FILE))) + help='The file the log will be written to') @click.option('--launch-tor/--no-launch-tor', - default=CTX.get('launch_tor', False), help='Launch Tor or try to connect to an existing Tor instance.') @click.option('--circuit-build-timeout', - default=CTX.get('circuit_build_timeout', 20), help='Option passed when launching Tor.') @click.version_option(BWSCAN_VERSION) @click.pass_context @@ -69,9 +61,12 @@ def cli(ctx, data_dir, loglevel, logfile, launch_tor, circuit_build_timeout): bandwidth measurements can then be aggregate to create the bandwidth values used by the Tor bandwidth authorities when creating the Tor consensus. """ + for k,v in ctx.default_map.items(): + if ctx.params.get(k) is None: + ctx.params[k] = v + # Create the data directory if it doesn't exist - data_dir = os.path.abspath(data_dir) - ctx.obj = ScanInstance(data_dir) + ctx.obj = ScanInstance(ctx.params.get('data_dir')) if not os.path.isdir(ctx.obj.measurement_dir): os.makedirs(ctx.obj.measurement_dir) @@ -81,7 +76,8 @@ def cli(ctx, data_dir, loglevel, logfile, launch_tor, circuit_build_timeout): ctx.obj.tor_dir) # Set up the logger to only output log lines of level `loglevel` and above. - setup_logging(log_level=loglevel, log_name=logfile) + setup_logging(log_level=ctx.params.get('loglevel'), + log_name=ctx.params.get('logfile')) @cli.command(short_help="Measure the Tor relays.") @@ -103,7 +99,8 @@ def scan(scan, partitions, current_partition, timeout, request_limit): # XXX: check that each run is producing the same input set! scan_time = str(int(time.time())) - scan_data_dir = os.path.join(scan.measurement_dir, '{}.running'.format(scan_time)) + scan_data_dir = os.path.join(scan.measurement_dir, + '{}.running'.format(scan_time)) if not os.path.isdir(scan_data_dir): os.makedirs(scan_data_dir) @@ -178,3 +175,8 @@ def aggregate(scan, scan_name, previous): scan.tor_state.addErrback(lambda failure: log.failure("Unexpected error")) scan.tor_state.addCallback(lambda _: reactor.stop()) reactor.run() + + +def start(): + config = read_config(os.path.join(DATA_DIR, CONFIG_FILE)) + return cli(default_map=config) diff --git a/setup.py b/setup.py index cf89041..df235e4 100644 --- a/setup.py +++ b/setup.py @@ -42,6 +42,6 @@ include_package_data=True, entry_points={ "console_scripts": [ - 'bwscan = bwscanner.scanner:cli', + 'bwscan = bwscanner.scanner:start', ]}, )