Change the way we get environment variables

Author: rw-access

detection_rules/eswrap.py

@@ -12,7 +12,7 @@
 from kibana import Kibana, RuleResource
 
 from .main import root
-from .misc import set_param_values
+from .misc import environ
 from .utils import normalize_timing_and_sort, unix_time_to_formatted, get_path
 from .rule_loader import get_rule, rta_mappings, load_rule_files, load_rules
 
@@ -179,10 +179,10 @@ def run(self, agent_hostname, indexes, verbose=True, **match):
 
 @es_group.command('collect-events')
 @click.argument('agent-hostname')
-@click.option('--elasticsearch-url', '-u', callback=set_param_values, expose_value=True)
-@click.option('--cloud-id', callback=set_param_values, expose_value=True)
-@click.option('--user', '-u', callback=set_param_values, expose_value=True, hide_input=False)
-@click.option('--password', '-p', callback=set_param_values, expose_value=True, hide_input=True)
+@click.option('--elasticsearch-url', '-u', default=environ("DR_ELASTICSEARCH_URL"))
+@click.option('--cloud-id', default=environ("DR_ELASTICSEARCH_URL"))
+@click.option('--user', '-u', default=environ("DR_USER"))
+@click.option('--password', '-p', default=environ("DR_PASSWORD"), hide_input=True)
 @click.option('--index', '-i', multiple=True, help='Index(es) to search against (default: all indexes)')
 @click.option('--agent-type', '-a', help='Restrict results to a source type (agent.type) ex: auditbeat')
 @click.option('--rta-name', '-r', help='Name of RTA in order to save events directly to unit tests data directory')
@@ -229,16 +229,23 @@ def normalize_file(events_file):
 
 @root.command("kibana-upload")
 @click.argument("toml-files", nargs=-1, required=True)
-@click.option('--kibana-url', '-u', callback=set_param_values, expose_value=True)
-@click.option('--cloud-id', callback=set_param_values, expose_value=True)
-@click.option('--user', '-u', callback=set_param_values, expose_value=True, hide_input=False)
-@click.option('--password', '-p', callback=set_param_values, expose_value=True, hide_input=True)
+@click.option('--kibana-url', '-u', default=environ("DR_KIBANA_URL"))
+@click.option('--cloud-id', default=environ("DR_CLOUD_ID"))
+@click.option('--user', '-u', default=environ("DR_USER"))
+@click.option('--password', '-p', default=environ("DR_PASSWORD"))
 def kibana_upload(toml_files, kibana_url, cloud_id, user, password):
     """Upload a list of rule .toml files to Kibana."""
     from uuid import uuid4
     from .packaging import manage_versions
     from .schemas import downgrade
 
+    if not cloud_id or kibana_url:
+        raise click.ClickException("Missing required --cloud-id or --kibana-url")
+
+    # don't prompt for these until there's a cloud id or kibana URL
+    user = user or click.prompt("user")
+    password = password or click.prompt("password", hide_input=True)
+
     with Kibana(cloud_id=cloud_id, url=kibana_url) as kibana:
         kibana.login(user, password)
 

detection_rules/misc.py

@@ -213,19 +213,6 @@ def parse_config():
     return _CONFIG
 
 
-def set_param_values(ctx, param, value):
-    """Get value for defined key."""
-    key = param.name
-    config = parse_config()
-    env_key = 'DR_' + key.upper()
-    prompt = True if param.hide_input is not False else False
-
-    if value:
-        return value
-    elif os.environ.get(env_key):
-        return os.environ[env_key]
-    elif config.get(key) is not None:
-        return config[key]
-    elif prompt:
-        return click.prompt(key, default=param.default if not param.default else None, hide_input=param.hide_input,
-                            show_default=True if param.default else False)
+def environ(name):
+    """Callback function for `default` to get an environment variable."""
+    return lambda: os.environ.get(name)