Merge pull request #22 from gormaniac/import-env-vars

feat: environment variables for user/pass in import

src/stormlibpp/_args.py

@@ -0,0 +1,16 @@
+"""CLI arguments shared by more than one script."""
+
+
+import argparse
+
+USER_PARSER = argparse.ArgumentParser(add_help=False)
+USER_PARSER.add_argument(
+    "--user",
+    help=(
+        "The Cortex user to authenticate with. The CORTEX_USER env var may be used"
+        " instead of this argument, however; this argument overrides the env var."
+        " If neither --user or CORTEX_USER are used, a prompt will appear to either"
+        " input a value or accept the default (return of getpass.getpass()). Only"
+        " works when connecting to a Cortex over HTTP."
+    ),
+)

src/stormlibpp/hstorm.py

@@ -8,9 +8,10 @@
 CLI commands are used, so the commands fail cleanly from the user's perspective.
 
 The script, like ``HttpCortex``, requires a user and password to communicate
-with the Synapse Cortex. A user can be passed via the command-line, otherwise
-``getpass`` is used to select the current user. A password is always prompted for.
-Future versions will allow for further user/pass customization options.
+with the Synapse Cortex. A user can be passed via the command-line, the ``CORTEX_USER``
+environment variable, or will be prompted for (``getpass`` is used as a default).
+The ``CORTEX_PASS`` environment variable may be used to give ``hstorm`` a password
+at runtime, otherwise one will be prompted for.
 
 The ``--no-verify`` option tells the script to not check the Cortex's HTTPS cert.
 This is needed to connect to any test Cortex or a Cortex that otherwise doesn't
@@ -20,33 +21,24 @@
 
 import asyncio
 import argparse
-import getpass
-import os
 import sys
 
+from ._args import USER_PARSER
 from .httpcore import HttpCortex
 from .output import OUTP
 from .stormcli import start_storm_cli
+from .utils import get_cortex_creds
 
 
 def get_args(argv: list[str]):
     """Build an argument parser for this script and parse the passed in args."""
 
-    args = argparse.ArgumentParser(prog="stormlibpp.hstorm")
+    args = argparse.ArgumentParser(prog="stormlibpp.hstorm", parents=[USER_PARSER,])
     args.add_argument("cortex", help="An HTTP URL for the Cortex.")
     args.add_argument("onecmd", nargs="?", help="A Storm command to run and exit.")
     args.add_argument(
         "-v", "--view", default=None, help="The iden of the Synapse View to use."
     )
-    args.add_argument(
-        "-u",
-        "--user",
-        default=None,
-        help=(
-            "The username to login to the Cortex with. "
-            "The value from getpass.getuser() is used if not given."
-        ),
-    )
     args.add_argument(
         "-n",
         "--no-verify",
@@ -68,19 +60,7 @@ async def main(argv: list[str]):
 
     args = get_args(argv)
 
-    if args.user:
-        username = args.user
-    elif (envusr := os.environ.get("CORTEX_USER")):
-        username = envusr
-    else:
-        gp_user = getpass.getuser()
-        user_in = input(f"Username [{gp_user}]: ")
-        username = user_in if user_in else gp_user
-
-    if (envpw := os.environ.get("CORTEX_PASS")):
-        password = envpw
-    else:
-        password = getpass.getpass()
+    username, password = get_cortex_creds(args.user)
 
     async with HttpCortex(args.cortex, username, password, ssl_verify=not args.no_verify) as hcore:
 

src/stormlibpp/import.py

@@ -19,9 +19,11 @@
 import synapse.telepath as s_telepath
 import yaml
 
+from ._args import USER_PARSER
 from .httpcore import HttpCortex
 from .output import handle_msg, log_storm_msg, OUTP
 from .stormcli import start_storm_cli
+from .utils import get_cortex_creds
 
 
 def endswith(text, items):
@@ -171,7 +173,7 @@ def cortex_proxy_contextmanager(cortex_url):
 
 
 def get_args(argv: list[str]):
-    parser = argparse.ArgumentParser(description=__doc__)
+    parser = argparse.ArgumentParser(description=__doc__, parents=[USER_PARSER])
     parser.add_argument(
         "folders", help="The folder(s) containing Storm scripts to execute", nargs="+"
     )
@@ -216,14 +218,6 @@ def get_args(argv: list[str]):
         help="Skips verification of the Cortex's certificate when using --http",
         action="store_true",
     )
-    parser.add_argument(
-        "--user",
-        help=(
-            "The Synapse user to authenticate with -"
-            " by default the return of getpass.getuser()"
-        ),
-        default=getpass.getuser(),
-    )
     parser.add_argument(
         "--view",
         help="An optional view to work in - otherwise the Cortex's default is chosen",
@@ -274,10 +268,9 @@ async def main(argv: list[str]):
         return "Can't use both --cortex and --local!"
     elif args.cortex:
         if args.http:
-            synuser = args.user
-            synpass = getpass.getpass()
+            username, password = get_cortex_creds(args.user)
             core_obj = functools.partial(
-                HttpCortex, args.cortex, synuser, synpass, ssl_verify=not args.no_verify
+                HttpCortex, args.cortex, username, password, ssl_verify=not args.no_verify
             )
         else:
             core_obj = functools.partial(cortex_proxy_contextmanager, args.cortex)

src/stormlibpp/utils.py

@@ -1,6 +1,7 @@
 """Miscellaneous functions that are helpful for working with Storm/Synapse."""
 
 
+import getpass
 import os
 
 import synapse.exc as s_exc
@@ -62,3 +63,21 @@ def normver(ver: str | tuple) -> tuple[str, tuple]:
         raise TypeError("Can only use a str or tuple as a Storm pkg version")
 
     return (verstr, vertup)
+
+
+def get_cortex_creds(_user):
+    if _user:
+        username = _user
+    elif (envusr := os.environ.get("CORTEX_USER")):
+        username = envusr
+    else:
+        gpusr = getpass.getuser()
+        inusr = input(f"Username [{gpusr}]: ")
+        username = inusr if inusr else gpusr
+
+    if (envpw := os.environ.get("CORTEX_PASS")):
+        password = envpw
+    else:
+        password = getpass.getpass()
+
+    return username, password