login-auth-tui 0.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,54 @@
1
+ Metadata-Version: 2.4
2
+ Name: login-auth-tui
3
+ Version: 0.1
4
+ Summary: TUI for login-auth stuff
5
+ Author: Nicholas Hurley
6
+ License-Expression: CC0-1.0
7
+ Requires-Dist: click
8
+ Requires-Dist: textual>=6.4.0
9
+ Requires-Python: >=3.13
10
+ Project-URL: CI, https://github.com/nwgh/login-auth-tui/actions
11
+ Project-URL: Changelog, https://github.com/nwgh/login-auth-tui/releases
12
+ Project-URL: Homepage, https://github.com/nwgh/login-auth-tui
13
+ Project-URL: Issues, https://github.com/nwgh/login-auth-tui/issues
14
+ Description-Content-Type: text/markdown
15
+
16
+ # login-auth-tui
17
+
18
+ [![Changelog](https://img.shields.io/github/v/release/nwgh/login-auth-tui?include_prereleases&label=changelog)](https://github.com/nwgh/login-auth-tui/releases)
19
+ [![Tests](https://github.com/nwgh/login-auth-tui/actions/workflows/test.yml/badge.svg)](https://github.com/nwgh/login-auth-tui/actions/workflows/test.yml)
20
+ [![License](https://img.shields.io/badge/license-CC0%201.0-blue.svg)](https://github.com/nwgh/login-auth-tui/blob/master/LICENSE)
21
+
22
+ TUI for login-auth stuff
23
+
24
+ ## Installation
25
+
26
+ Install this tool from the releases page
27
+
28
+ ## Usage
29
+
30
+ For help, run:
31
+ ```bash
32
+ login-auth-tui --help
33
+ ```
34
+
35
+ ## Development
36
+
37
+ To contribute to this tool, first install `uv`. See [the uv documentation](http
38
+ s://docs.astral.sh/uv/getting-started/installation/) for how.
39
+
40
+ Next, checkout the code
41
+ ```bash
42
+ git clone https://github.com/nwgh/login-auth-tui
43
+ ```
44
+
45
+ Then create a new virtual environment and sync the dependencies:
46
+ ```bash
47
+ cd login-auth-tui
48
+ uv sync
49
+ ```
50
+
51
+ To run the tests:
52
+ ```bash
53
+ uv run python -m pytest
54
+ ```
@@ -0,0 +1,39 @@
1
+ # login-auth-tui
2
+
3
+ [![Changelog](https://img.shields.io/github/v/release/nwgh/login-auth-tui?include_prereleases&label=changelog)](https://github.com/nwgh/login-auth-tui/releases)
4
+ [![Tests](https://github.com/nwgh/login-auth-tui/actions/workflows/test.yml/badge.svg)](https://github.com/nwgh/login-auth-tui/actions/workflows/test.yml)
5
+ [![License](https://img.shields.io/badge/license-CC0%201.0-blue.svg)](https://github.com/nwgh/login-auth-tui/blob/master/LICENSE)
6
+
7
+ TUI for login-auth stuff
8
+
9
+ ## Installation
10
+
11
+ Install this tool from the releases page
12
+
13
+ ## Usage
14
+
15
+ For help, run:
16
+ ```bash
17
+ login-auth-tui --help
18
+ ```
19
+
20
+ ## Development
21
+
22
+ To contribute to this tool, first install `uv`. See [the uv documentation](http
23
+ s://docs.astral.sh/uv/getting-started/installation/) for how.
24
+
25
+ Next, checkout the code
26
+ ```bash
27
+ git clone https://github.com/nwgh/login-auth-tui
28
+ ```
29
+
30
+ Then create a new virtual environment and sync the dependencies:
31
+ ```bash
32
+ cd login-auth-tui
33
+ uv sync
34
+ ```
35
+
36
+ To run the tests:
37
+ ```bash
38
+ uv run python -m pytest
39
+ ```
@@ -0,0 +1,71 @@
1
+ [project]
2
+ name = "login-auth-tui"
3
+ version = "0.1"
4
+ description = "TUI for login-auth stuff"
5
+ readme = "README.md"
6
+ authors = [{ name = "Nicholas Hurley" }]
7
+ license = "CC0-1.0"
8
+ requires-python = ">=3.13"
9
+ classifiers = []
10
+ dependencies = ["click", "textual>=6.4.0"]
11
+
12
+ [build-system]
13
+ requires = ["uv_build>=0.9.5,<0.10.0"]
14
+ build-backend = "uv_build"
15
+
16
+ [project.urls]
17
+ Homepage = "https://github.com/nwgh/login-auth-tui"
18
+ Changelog = "https://github.com/nwgh/login-auth-tui/releases"
19
+ Issues = "https://github.com/nwgh/login-auth-tui/issues"
20
+ CI = "https://github.com/nwgh/login-auth-tui/actions"
21
+
22
+ [project.scripts]
23
+ auth-manage = "login_auth_tui.cli:cli"
24
+
25
+ [dependency-groups]
26
+ dev = [
27
+ { include-group = "autoflake" },
28
+ { include-group = "bandit" },
29
+ { include-group = "black" },
30
+ { include-group = "flake8" },
31
+ { include-group = "installer" },
32
+ { include-group = "isort" },
33
+ { include-group = "test" },
34
+ "pre-commit",
35
+ "textual-dev",
36
+ ]
37
+ autoflake = ["autoflake"]
38
+ bandit = ["bandit", "tomli"]
39
+ black = ["black"]
40
+ flake8 = ["flake8", "flake8-bugbear", "Flake8-pyproject", "pep8-naming"]
41
+ isort = ["isort"]
42
+ installer = ["pyinstaller"]
43
+ test = ["pytest"]
44
+
45
+ [tool.black]
46
+ line-length = 99
47
+
48
+ [tool.isort]
49
+ line_length = 99
50
+ profile = "black"
51
+
52
+ [tool.pytest.ini_options]
53
+ console_output_style = "progress"
54
+ minversion = "7.2.2"
55
+ python_files = "test_*.py"
56
+ addopts = "-vvvv --ignore-glob .venv* --ignore-glob venv*"
57
+
58
+ [tool.flake8]
59
+ max-line-length = 99
60
+ exclude = [".venv*", "venv*", "typings"]
61
+ extend-ignore = ["E203", "E501"]
62
+ extend-select = ["B950"]
63
+
64
+ [tool.bandit]
65
+ exclude_dirs = [".venv", "venv"]
66
+
67
+ [[tool.uv.index]]
68
+ name = "testpypi"
69
+ url = "https://test.pypi.org/simple/"
70
+ publish-url = "https://test.pypi.org/legacy/"
71
+ explicit = true
File without changes
@@ -0,0 +1,4 @@
1
+ from login_auth_tui.cli import cli
2
+
3
+ if __name__ == "__main__":
4
+ cli()
@@ -0,0 +1,61 @@
1
+ import configparser
2
+ import datetime
3
+ import logging
4
+ import os
5
+
6
+ from .util import log_command, log_command_output
7
+
8
+ logger = logging.getLogger("aws")
9
+
10
+ HOMEDIR = str(os.getenv("HOME"))
11
+
12
+
13
+ def credential_is_valid(profile: str) -> bool:
14
+ now = datetime.datetime.now(datetime.UTC)
15
+
16
+ p = configparser.ConfigParser()
17
+ p.read(os.path.join(HOMEDIR, ".aws", "credentials"))
18
+
19
+ then_s = p.get(profile, "expiration")
20
+ then = datetime.datetime.strptime(f"{then_s} +0000", "%Y-%m-%d %H:%M:%S %z")
21
+
22
+ if now < then:
23
+ return True
24
+ return False
25
+
26
+
27
+ def aws_mfa(profile: str, force: bool) -> None:
28
+ logger.info(f"Running mfa process for {profile}")
29
+
30
+ no_mfa_file = os.path.join(HOMEDIR, f".no-mfa-{profile}")
31
+ if os.path.exists(no_mfa_file) and not force:
32
+ logger.info("Not doing mfa due to temp file existing. (Removing the file.)")
33
+ os.unlink(no_mfa_file)
34
+ return
35
+
36
+ if credential_is_valid(profile) and not force:
37
+ logger.info("Not doing mfa due to credentials still valid.")
38
+ return
39
+
40
+ code = log_command_output(
41
+ logger, ["op", "read", "op://Private/AWS/one-time password?attribute=totp"]
42
+ )
43
+ logger.info(f"Using {code} for {profile}")
44
+ rval = log_command(
45
+ logger,
46
+ [f"{os.getenv('HOMEBREW_PREFIX')}/bin/aws", "mfa", profile],
47
+ input=code.encode("utf-8"),
48
+ ).returncode
49
+
50
+ if rval != 0:
51
+ # For some reason, recently, aws-mfa has been returning 1 even though
52
+ # it successfully updates the tokens. So if it does return non-zero,
53
+ # we need to check the credentials file itself to verify if it's a real
54
+ # failure or not.
55
+ logger.info("Warning: MFA seems to have failed by rval, checking expiration")
56
+ if not credential_is_valid(profile):
57
+ logger.error(f"Not rotating long-term keys due to MFA failure ({rval})")
58
+ raise Exception("AWS MFA failure")
59
+
60
+ logger.info("Rotating long-term keys")
61
+ log_command(logger, ["aws", "rotate-iam-keys", f"{profile}-long-term"], check=True)
@@ -0,0 +1,253 @@
1
+ import click
2
+ import copy
3
+ import logging
4
+ import os
5
+ import sqlite3
6
+ import sys
7
+ import time
8
+
9
+ from .aws_mfa import aws_mfa
10
+ from .dbx_token_rotate import manage_dbx_tokens
11
+ from .kion_auth_password import update_kion_auth_password
12
+ from .kion_cli_password import DEFAULT_CONFIG_FILE, update_kion_cli_password
13
+ from .tui import run_tui
14
+ from .tui_bootstrap import bootstrap_tui_config
15
+ from .util import configure_logging, log_command
16
+
17
+ TWELVE_HOURS = 12 * 60 * 60
18
+ DEFAULT_LOG = os.path.join(str(os.getenv("HOME")), ".logs", "login-auth.log")
19
+
20
+
21
+ class CliArgs:
22
+ __slots__ = ["force", "log", "level"]
23
+
24
+ def __init__(self):
25
+ self.force = False
26
+ self.log = DEFAULT_LOG
27
+ self.level = "INFO"
28
+
29
+
30
+ @click.group()
31
+ @click.version_option()
32
+ @click.option(
33
+ "--force",
34
+ is_flag=True,
35
+ default=False,
36
+ help="Force update regardless of last update time (only aws, batch, and dbx).",
37
+ )
38
+ @click.option(
39
+ "--log",
40
+ type=click.Path(dir_okay=False, writable=True, resolve_path=True, allow_dash=True),
41
+ default=DEFAULT_LOG,
42
+ help="File to log to ('-' for stdout).",
43
+ )
44
+ @click.option(
45
+ "--level",
46
+ type=click.Choice([k for k in logging.getLevelNamesMapping().keys() if k != "NOTSET"]),
47
+ default="INFO",
48
+ help="Logging level to use.",
49
+ )
50
+ @click.pass_context
51
+ def cli(ctx: click.Context, force: bool, log: str, level: str):
52
+ """Run authentication useful when developing on DataConnect."""
53
+ ctx.ensure_object(CliArgs)
54
+ ctx.obj.force = force
55
+ ctx.obj.log = log
56
+ ctx.obj.level = level
57
+
58
+
59
+ def login_auth(force: bool, conn: sqlite3.Connection) -> None:
60
+ # Now that we hold an exclusive lock on the config, see if we actually need
61
+ # to do anything.
62
+ last_update = conn.execute("SELECT * FROM last_update").fetchone()[0]
63
+ now = time.time()
64
+ if last_update < (now - TWELVE_HOURS):
65
+ logging.info("More than twelve hours since last update.")
66
+ elif force:
67
+ logging.info("Running due to force")
68
+ else:
69
+ logging.info("Not running; auth does not need to be updated")
70
+ return
71
+
72
+ res = conn.execute("SELECT * FROM config")
73
+ config = dict(zip([r[0] for r in res.description], res.fetchone()))
74
+
75
+ if config["wait"]:
76
+ logging.info(f"Waiting {config['wait']} seconds for boot...")
77
+ time.sleep(config["wait"])
78
+
79
+ if config["ssh_add"]:
80
+ logging.info("Adding SSH keys...")
81
+ env = copy.deepcopy(os.environ)
82
+ args = ["/usr/bin/ssh-add"]
83
+ if sys.platform == "darwin":
84
+ env["APPLE_SSH_ADD_BEHAVIOR"] = "yes"
85
+ args.append("-A")
86
+ log_command(logging.getLogger(), args, env=env)
87
+
88
+ success = True
89
+ res = conn.execute("SELECT * FROM aws_profiles")
90
+ aws_profiles = [r[0] for r in res.fetchall()]
91
+ for aws_profile in aws_profiles:
92
+ logging.info(f"Authenticating AWS profile {aws_profile}")
93
+ try:
94
+ aws_mfa(aws_profile, force)
95
+ except Exception:
96
+ logging.exception(f"Exception running AWS MFA {aws_profile}")
97
+ success = False
98
+
99
+ res = conn.execute("SELECT * FROM dbx_profiles")
100
+ for row in res.fetchall():
101
+ profile, alias = row
102
+ logging.info(f"Rotating DBX token {profile} (alias={alias})")
103
+ try:
104
+ manage_dbx_tokens(profile, [alias], force)
105
+ except Exception:
106
+ logging.exception(f"Exception rotating DBX token {profile}")
107
+ success = False
108
+
109
+ kion_auth_config_file = config["kion_auth_config_file"]
110
+ if kion_auth_config_file:
111
+ logging.info("Update kion-auth config file with current password")
112
+ update_kion_auth_password(kion_auth_config_file)
113
+
114
+ kion_cli_config_file = config["kion_cli_config_file"]
115
+ if kion_cli_config_file:
116
+ logging.info("Update kion cli config file with current password")
117
+ update_kion_cli_password(kion_cli_config_file)
118
+
119
+ if success:
120
+ logging.info("Success!")
121
+ conn.execute(f"UPDATE last_update SET tstamp = {int(now)}") # nosec
122
+ else:
123
+ logging.info("Failed")
124
+
125
+
126
+ @cli.command("batch")
127
+ @click.option("--background", is_flag=True, default=False, help="Run in the background.")
128
+ @click.option(
129
+ "--config",
130
+ type=click.Path(dir_okay=False, exists=True, resolve_path=True),
131
+ default=os.path.join(str(os.getenv("HOME")), ".config", "login-auth.sqlite"),
132
+ help="Path to configuration database.",
133
+ )
134
+ @click.pass_context
135
+ def login_auth_main(ctx: click.Context, background: bool, config: str) -> None:
136
+ """Run all the login management in one batch."""
137
+ if background:
138
+ # TODO - run in the background
139
+ pid = os.fork()
140
+ if pid != 0:
141
+ # This is the parent, we done
142
+ return
143
+
144
+ # This is the child, detach from the parent
145
+ os.setpgrp()
146
+
147
+ configure_logging(ctx.obj.log, ctx.obj.level)
148
+
149
+ conn = sqlite3.connect(config)
150
+ conn.execute("BEGIN EXCLUSIVE TRANSACTION")
151
+
152
+ try:
153
+ login_auth(ctx.obj.force, conn)
154
+ except Exception:
155
+ logging.exception("Exception doing login_auth")
156
+ conn.execute("ROLLBACK")
157
+ raise
158
+ else:
159
+ logging.info("Commit transaction")
160
+ conn.execute("COMMIT")
161
+
162
+
163
+ @cli.command("aws")
164
+ @click.argument("profile")
165
+ @click.pass_context
166
+ def aws_mfa_main(ctx: click.Context, profile: str) -> None:
167
+ """Run AWS MFA process.
168
+
169
+ PROFILE the AWS profile to run the MFA process for.
170
+ """
171
+ configure_logging(ctx.obj.log, ctx.obj.level)
172
+
173
+ aws_mfa(profile, ctx.obj.force)
174
+
175
+
176
+ @cli.command("dbx")
177
+ @click.option("--alias", multiple=True, help="Other profiles that are aliases.")
178
+ @click.argument("profile")
179
+ @click.pass_context
180
+ def dbx_token_rotate_main(ctx: click.Context, alias: tuple[str], profile: str) -> None:
181
+ """Rotate Databricks API token.
182
+
183
+ PROFILE the Databricks configuration profile to rotate the API token for.
184
+ """
185
+ configure_logging(ctx.obj.log, ctx.obj.level)
186
+
187
+ manage_dbx_tokens(profile, alias, ctx.obj.force)
188
+
189
+
190
+ @cli.command("kion-auth")
191
+ @click.option(
192
+ "--file",
193
+ type=click.Path(exists=True, dir_okay=False, writable=True, resolve_path=True),
194
+ default=os.path.join(str(os.getenv("HOME")), ".config", "kion-auth.ini"),
195
+ help="Path of kion-auth configuration file.",
196
+ )
197
+ @click.pass_context
198
+ def update_kion_auth_password_main(ctx: click.Context, file: str) -> None:
199
+ """Update the password used by kion-auth from 1Password."""
200
+ configure_logging(ctx.obj.log, ctx.obj.level)
201
+
202
+ update_kion_auth_password(file)
203
+
204
+
205
+ @cli.command("kion-cli")
206
+ @click.option(
207
+ "--file",
208
+ type=click.Path(exists=True, dir_okay=False, writable=True, resolve_path=True),
209
+ default=DEFAULT_CONFIG_FILE,
210
+ help="Path of kion cli configuration file.",
211
+ )
212
+ @click.pass_context
213
+ def update_kion_cli_password_main(ctx: click.Context, file: str) -> None:
214
+ """Update the password used by the kion cli from 1Password."""
215
+ configure_logging(ctx.obj.log, ctx.obj.level)
216
+
217
+ update_kion_cli_password(file)
218
+
219
+
220
+ @cli.command("tui")
221
+ @click.option(
222
+ "--config",
223
+ type=click.Path(exists=True, dir_okay=False, writable=True, resolve_path=True),
224
+ default=os.path.join(str(os.getenv("HOME")), ".config", "login_auth.yml"),
225
+ help="Path of TUI configuration file.",
226
+ )
227
+ @click.pass_context
228
+ def tui(ctx: click.Context, config: str) -> None:
229
+ """Run the TUI for authentication processes."""
230
+ configure_logging(ctx.obj.log, ctx.obj.level)
231
+ run_tui(config)
232
+
233
+
234
+ @cli.command("tui-bootstrap")
235
+ @click.option(
236
+ "--in-config",
237
+ type=click.Path(dir_okay=False, exists=True, resolve_path=True),
238
+ default=os.path.join(str(os.getenv("HOME")), ".config", "login-auth.sqlite"),
239
+ help="Path to configuration database.",
240
+ )
241
+ @click.option(
242
+ "--out-config",
243
+ type=click.Path(dir_okay=False, writable=True, resolve_path=True),
244
+ default=os.path.join(str(os.getenv("HOME")), ".config", "login_auth.yml"),
245
+ help="Path of TUI configuration file.",
246
+ )
247
+ @click.pass_context
248
+ def tui_bootstrap(ctx: click.Context, in_config: str, out_config: str) -> None:
249
+ """Bootstrap the tui configuration from the batch configuration."""
250
+ configure_logging(ctx.obj.log, ctx.obj.level)
251
+ if os.path.exists(out_config) and not ctx.obj.force:
252
+ raise ValueError(f"Not overwriting existing TUI config {out_config}")
253
+ bootstrap_tui_config(in_config, out_config)
@@ -0,0 +1,90 @@
1
+ import configparser
2
+ import datetime
3
+ import json
4
+ import logging
5
+ import os
6
+ import shutil
7
+ import time
8
+ import typing
9
+
10
+ from .util import log_command_output
11
+
12
+ logger = logging.getLogger("dbx")
13
+
14
+ HOMEDIR = os.getenv("HOME")
15
+ TEN_DAYS_MS = 10 * 24 * 60 * 60 * 1000
16
+ NINETY_DAYS_S = int((TEN_DAYS_MS / 1000) * 9)
17
+
18
+
19
+ def run_dbx(profile: str, *args) -> str:
20
+ subprocess_args = ["databricks", "--profile", profile, "-o", "json"] + list(args)
21
+ logger.info(f"Run {' '.join(subprocess_args)}")
22
+ return log_command_output(logger, subprocess_args)
23
+
24
+
25
+ def create_new_token(profile: str) -> str:
26
+ logger.info(f"Creating new token for {profile}")
27
+ today = str(datetime.date.today())
28
+ new_token_json = run_dbx(
29
+ profile,
30
+ "tokens",
31
+ "create",
32
+ "--comment",
33
+ f"dbx cli {today}",
34
+ "--lifetime-seconds",
35
+ str(NINETY_DAYS_S),
36
+ )
37
+ new_token = json.loads(new_token_json)
38
+
39
+ return new_token["token_value"]
40
+
41
+
42
+ def delete_old_token(profile: str, token_id: str) -> None:
43
+ logger.info(f"Deleting old token {token_id} for profile {profile}")
44
+ run_dbx(profile, "tokens", "delete", token_id)
45
+
46
+
47
+ def overwrite_dbx_config(
48
+ dbxcfg: str,
49
+ cp: configparser.ConfigParser,
50
+ profile: str,
51
+ aliases: typing.Iterable[str],
52
+ new_token: str,
53
+ ) -> None:
54
+ logger.info(f"Updating config file for profile {profile} (alias={aliases})")
55
+ cp[profile]["token"] = new_token
56
+ cp[profile]["nwgh-last-update"] = str(datetime.date.today())
57
+ if aliases:
58
+ for alias in aliases:
59
+ cp[alias]["token"] = new_token
60
+ with open(f"{dbxcfg}.new", "w") as f:
61
+ cp.write(f)
62
+ shutil.copyfile(dbxcfg, f"{dbxcfg}.old")
63
+ shutil.move(f"{dbxcfg}.new", dbxcfg)
64
+
65
+
66
+ def manage_dbx_tokens(profile: str, aliases: typing.Iterable[str], force: bool) -> None:
67
+ dbxcfg = os.path.join(str(os.getenv("HOME")), ".databrickscfg")
68
+ if not os.path.exists(dbxcfg):
69
+ raise Exception(f"Databricks config file {dbxcfg} does not exist")
70
+
71
+ cp = configparser.ConfigParser()
72
+ cp.read(dbxcfg)
73
+ if profile not in cp:
74
+ raise Exception(f"Databricks config does not have profile {dbxcfg}")
75
+
76
+ json_result = run_dbx(profile, "tokens", "list")
77
+ tokens = json.loads(json_result)
78
+ if len(tokens) != 1:
79
+ raise Exception(
80
+ f"This script does not work with {len(tokens)} tokens! "
81
+ "There can be only one."
82
+ )
83
+
84
+ now_ms = int(time.time()) * 1000
85
+ if force or (tokens[0]["expiry_time"] - now_ms <= TEN_DAYS_MS):
86
+ new_token = create_new_token(profile)
87
+ overwrite_dbx_config(dbxcfg, cp, profile, aliases, new_token)
88
+ delete_old_token(profile, tokens[0]["token_id"])
89
+ else:
90
+ logger.info(f"Token for {profile} is still ok, not rotating")
@@ -0,0 +1,28 @@
1
+ import configparser
2
+ import logging
3
+ import os
4
+
5
+ from .util import log_command_output
6
+
7
+ logger = logging.getLogger("kion-auth")
8
+
9
+
10
+ def update_kion_auth_password(config_file: str) -> None:
11
+ if not os.path.exists(config_file):
12
+ logger.error(f"kion-auth config file {config_file} does not exist")
13
+ return
14
+
15
+ cp = configparser.ConfigParser()
16
+ cp.read(config_file)
17
+
18
+ if not cp.get("kion_auth", "password"):
19
+ logger.info("Kion password is not defined, not updating")
20
+ return
21
+
22
+ password = log_command_output(
23
+ logger, ["op", "read", "op://Private/CMS EUA/password"]
24
+ )
25
+ cp["kion_auth"]["password"] = password
26
+
27
+ with open(config_file, "w") as f:
28
+ cp.write(f)
@@ -0,0 +1,29 @@
1
+ import logging
2
+ import os
3
+
4
+ from .util import log_command_output
5
+
6
+ logger = logging.getLogger("kion-cli")
7
+
8
+ DEFAULT_CONFIG_FILE = os.path.join(str(os.getenv("HOME")), ".config", "kion.yml")
9
+
10
+
11
+ def update_kion_cli_password(config_file: str) -> None:
12
+ if not os.path.exists(config_file):
13
+ logger.error(f"kion cli config file {config_file} does not exist")
14
+ return
15
+
16
+ password = log_command_output(
17
+ logger, ["op", "read", "op://Private/CMS EUA/password"]
18
+ )
19
+
20
+ new_lines = []
21
+ with open(config_file) as f:
22
+ for line in f:
23
+ line = line.rstrip()
24
+ if line.startswith(" password:"):
25
+ line = f" password: {password}"
26
+ new_lines.append(line)
27
+
28
+ with open(config_file, "w") as f:
29
+ f.write("\n".join(new_lines))
@@ -0,0 +1,17 @@
1
+ import logging
2
+ import random
3
+ import datetime
4
+
5
+ logger = logging.getLogger("noop")
6
+
7
+
8
+ def run_noop():
9
+ now = datetime.datetime.now()
10
+ nlines = random.randint(2, 5)
11
+ for i in range(1, nlines):
12
+ logger.info(f"Message {i} from run at {now}")
13
+
14
+ if random.choice([0, 1]):
15
+ msg = f"Faking error for run at {now}"
16
+ logger.error(msg)
17
+ raise Exception(msg)
@@ -0,0 +1,259 @@
1
+ import enum
2
+ import io
3
+ import logging
4
+ from textual.screen import ModalScreen
5
+ import yaml
6
+
7
+ from textual import on, work
8
+ from textual.app import App, ComposeResult
9
+ from textual.binding import Binding
10
+ from textual.containers import HorizontalGroup
11
+ from textual.logging import TextualHandler
12
+ from textual.message import Message
13
+ from textual.reactive import reactive
14
+ from textual.widget import Widget
15
+ from textual.widgets import (
16
+ Footer,
17
+ Label,
18
+ ListItem,
19
+ ListView,
20
+ Log,
21
+ Rule,
22
+ )
23
+
24
+ from .aws_mfa import aws_mfa
25
+ from .dbx_token_rotate import manage_dbx_tokens
26
+ from .kion_cli_password import DEFAULT_CONFIG_FILE, update_kion_cli_password
27
+ from .noop import run_noop
28
+
29
+ logger = logging.getLogger("tui")
30
+
31
+
32
+ class ProcessStatus(enum.Enum):
33
+ OK = 0
34
+ FAIL = 1
35
+ UNKNOWN = 2
36
+ RUNNING = 3
37
+
38
+
39
+ # Each process will write to its own io.StringIO
40
+ class AuthProcess:
41
+ def __init__(self, process_config: dict):
42
+ self.name = process_config["name"]
43
+ self.type = process_config["type"]
44
+
45
+ if self.type == "aws":
46
+ self.profile = process_config["profile"]
47
+ elif self.type == "dbx":
48
+ self.profile = process_config["profile"]
49
+ self.aliases = []
50
+ alias = process_config.get("alias")
51
+ if alias:
52
+ self.aliases.append(alias)
53
+ elif self.type == "kion-cli":
54
+ self.config = process_config.get("config", DEFAULT_CONFIG_FILE)
55
+ elif self.type == "noop":
56
+ # This is used for testing things
57
+ pass
58
+ else:
59
+ raise ValueError(f"Unexpected auth process type: {self.type}")
60
+
61
+ self.log_stream = io.StringIO()
62
+ self.log_handler = logging.StreamHandler(self.log_stream)
63
+
64
+ def run(self):
65
+ logger.debug(f"Run {self.name} {self.type}")
66
+ process_logger = logging.getLogger(self.type)
67
+ process_logger.addHandler(self.log_handler)
68
+ try:
69
+ if self.type == "aws":
70
+ aws_mfa(self.profile, False)
71
+ elif self.type == "dbx":
72
+ manage_dbx_tokens(self.profile, self.aliases, False)
73
+ elif self.type == "kion-cli":
74
+ update_kion_cli_password(self.config)
75
+ elif self.type == "noop":
76
+ run_noop()
77
+ else:
78
+ raise ValueError(f"Unexpected auth process type: {self.type}")
79
+ finally:
80
+ process_logger.removeHandler(self.log_handler)
81
+
82
+
83
+ def load_config(config_file: str) -> list[AuthProcess]:
84
+ logger.debug(f"Loading config from {config_file}")
85
+ processes = []
86
+ with open(config_file) as f:
87
+ config = yaml.safe_load(f)
88
+ for auth_process in config:
89
+ processes.append(AuthProcess(auth_process))
90
+ return processes
91
+
92
+
93
+ class AuthProcessStatus(Widget):
94
+ value = reactive(ProcessStatus.UNKNOWN)
95
+
96
+ def render(self) -> str:
97
+ if self.value == ProcessStatus.OK:
98
+ return "\u2714"
99
+ if self.value == ProcessStatus.FAIL:
100
+ return "\u2716"
101
+ if self.value == ProcessStatus.RUNNING:
102
+ return "-"
103
+ return "?"
104
+
105
+
106
+ class AuthRunMessage(Message):
107
+ def __init__(self, status: ProcessStatus, *args, **kwargs):
108
+ super().__init__(*args, **kwargs)
109
+ self.status = status
110
+
111
+
112
+ class AuthRunCompleteMessage(Message):
113
+ pass
114
+
115
+
116
+ class AuthProcessItem(ListItem):
117
+ def __init__(self, proc: AuthProcess, *args, **kwargs):
118
+ super().__init__(*args, **kwargs)
119
+ self.process_config = proc
120
+ self.status = AuthProcessStatus()
121
+
122
+ def compose(self) -> ComposeResult:
123
+ # TODO - make use of more stuff here
124
+ yield HorizontalGroup(Label(self.process_config.name), Label(" "), self.status)
125
+
126
+ @on(AuthRunMessage)
127
+ def run_complete(self, msg: AuthRunMessage):
128
+ self.status.value = msg.status
129
+ self.post_message(AuthRunCompleteMessage())
130
+
131
+ @work(exclusive=True, thread=True)
132
+ def run_auth(self):
133
+ logger.debug(f"Run auth: {self}")
134
+ result = ProcessStatus.UNKNOWN
135
+ try:
136
+ self.process_config.run()
137
+ result = ProcessStatus.OK
138
+ except Exception:
139
+ logger.exception("Error updating auth")
140
+ result = ProcessStatus.FAIL
141
+ self.post_message(AuthRunMessage(result))
142
+
143
+
144
+ class UpdateLogMessage(Message):
145
+ def __init__(self, name: str, value: str, *args, **kwargs):
146
+ super().__init__(*args, **kwargs)
147
+ self.proc_name = name
148
+ self.log_value = value
149
+
150
+
151
+ class AuthProcessList(ListView):
152
+ BINDINGS = [
153
+ Binding("j, down", "cursor_down", "Cursor down"),
154
+ Binding("k, up", "cursor_up", "Cursor up"),
155
+ Binding("o, enter", "select_cursor", "Select"),
156
+ Binding("r", "run_auth", "Run highlighted"),
157
+ ]
158
+
159
+ def __init__(self, procs: list[AuthProcess], *args, **kwargs):
160
+ super().__init__(*args, **kwargs)
161
+ self._auth_processes: list[AuthProcess] = procs
162
+ self._highlighted: AuthProcessItem
163
+ self.running = False
164
+ self.running_all = False
165
+ self.current_item = 0
166
+ self.set_interval(300, self.run_all)
167
+
168
+ def compose(self) -> ComposeResult:
169
+ for p in self._auth_processes:
170
+ yield AuthProcessItem(p)
171
+
172
+ def action_run_auth(self):
173
+ self.running = True
174
+ self._highlighted.status.value = ProcessStatus.RUNNING
175
+ self._highlighted.run_auth()
176
+
177
+ def on_list_view_selected(self, event: ListView.Selected):
178
+ logger.debug(f"Opening {event.item}")
179
+ proc: AuthProcessItem = event.item # type: ignore
180
+ proc_name = proc.process_config.name
181
+ log_value = proc.process_config.log_stream.getvalue()
182
+ self.post_message(UpdateLogMessage(proc_name, log_value))
183
+
184
+ def on_list_view_highlighted(self, event: ListView.Highlighted):
185
+ logger.debug(f"Item is {event.item}")
186
+ self._highlighted = event.item # type: ignore
187
+
188
+ def run_all(self):
189
+ self.running_all = True
190
+ self.running = True
191
+ self.current_item = 0
192
+ self.run_next()
193
+
194
+ def run_next(self):
195
+ if self.current_item >= len(self.children):
196
+ self.running_all = False
197
+ else:
198
+ self.running = True
199
+ item: AuthProcessItem = self.children[self.current_item] # type: ignore
200
+ item.run_auth()
201
+
202
+ @on(AuthRunCompleteMessage)
203
+ def handle_auth_run_complete(self):
204
+ self.running = False
205
+ if self.running_all:
206
+ self.current_item += 1
207
+ self.run_next()
208
+
209
+
210
+ class ProcessLogHeader(Widget):
211
+ process = reactive("none")
212
+
213
+ def render(self):
214
+ return f"Log for process: {self.process}"
215
+
216
+
217
+ class LogScreen(ModalScreen):
218
+ BINDINGS = [Binding("x", "close_log", "Close")]
219
+
220
+ def __init__(self, auth_name: str, log_value: str, *args, **kwargs):
221
+ super().__init__(*args, **kwargs)
222
+ self.auth_name = auth_name
223
+ self.log_value = log_value
224
+
225
+ def compose(self) -> ComposeResult:
226
+ yield Label(f"Log for {self.auth_name}")
227
+ yield Rule()
228
+
229
+ log = Log()
230
+ yield log
231
+ log.write(self.log_value)
232
+ yield Footer()
233
+
234
+ def action_close_log(self):
235
+ self.app.pop_screen()
236
+
237
+
238
+ class AuthTUI(App):
239
+ BINDINGS = [Binding("q", "quit", "Quit")]
240
+ CSS_PATH = "tui.tcss"
241
+
242
+ def __init__(self, processes: list[AuthProcess], **kwargs):
243
+ super().__init__(**kwargs)
244
+ self._processes = processes
245
+
246
+ def compose(self) -> ComposeResult:
247
+ yield AuthProcessList(self._processes, id="auth_list")
248
+ yield Footer()
249
+
250
+ @on(UpdateLogMessage)
251
+ def handle_update_log(self, msg: UpdateLogMessage):
252
+ self.push_screen(LogScreen(msg.proc_name, msg.log_value))
253
+
254
+
255
+ def run_tui(config_file: str):
256
+ logger.addHandler(TextualHandler())
257
+ processes = load_config(config_file)
258
+ app = AuthTUI(processes)
259
+ app.run()
@@ -0,0 +1,11 @@
1
+ ProcessLogHeader {
2
+ height: 1;
3
+ }
4
+
5
+ LogScreen {
6
+ align: center middle;
7
+ height: 80%;
8
+ width: 80%;
9
+ background: $surface;
10
+ padding: 0 1;
11
+ }
@@ -0,0 +1,33 @@
1
+ import sqlite3
2
+ import yaml
3
+
4
+
5
+ def bootstrap_tui_config(input_config: str, output_config: str):
6
+ tui_config = []
7
+ with sqlite3.connect(input_config) as conn:
8
+ res = conn.execute("SELECT * FROM aws_profiles").fetchall()
9
+ for r in res:
10
+ profile_name = r[0]
11
+ tui_config.append(
12
+ {"name": f"AWS ({profile_name})", "type": "aws", "profile": profile_name}
13
+ )
14
+
15
+ res = conn.execute("SELECT * FROM dbx_profiles").fetchall()
16
+ for r in res:
17
+ profile_name = r[0]
18
+ alias = r[1]
19
+ config_entry = {
20
+ "name": f"Databricks ({profile_name})",
21
+ "type": "dbx",
22
+ "profile": profile_name,
23
+ }
24
+ if alias:
25
+ config_entry["alias"] = alias
26
+ tui_config.append(config_entry)
27
+
28
+ res = conn.execute("SELECT kion_cli_config_file FROM config").fetchall()
29
+ if res[0][0]:
30
+ tui_config.append({"name": "Kion", "type": "kion-cli", "config": res[0][0]})
31
+
32
+ with open(output_config, "w") as f:
33
+ yaml.dump(tui_config, f)
@@ -0,0 +1,50 @@
1
+ import logging
2
+ import logging.handlers
3
+ import os
4
+ import subprocess # nosec
5
+ import sys
6
+
7
+
8
+ HOME = str(os.getenv("HOME"))
9
+
10
+
11
+ def configure_logging(log: str, level: str) -> None:
12
+ if log in ("stdout", "-"):
13
+ handler = logging.StreamHandler(sys.stdout)
14
+ else:
15
+ handler = logging.handlers.RotatingFileHandler(
16
+ log,
17
+ backupCount=1,
18
+ maxBytes=1024 * 1024,
19
+ )
20
+ try:
21
+ log_level = logging.getLevelNamesMapping()[level]
22
+ except Exception:
23
+ raise ValueError(f"Unhandled log level {level}")
24
+ logging.basicConfig(
25
+ format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
26
+ level=log_level,
27
+ handlers=[handler],
28
+ )
29
+
30
+
31
+ def log_command(
32
+ logger: logging.Logger, command: list[str], **kwargs
33
+ ) -> subprocess.CompletedProcess:
34
+ logger.info(f"Run {' '.join(command)}")
35
+ result = subprocess.run(command, capture_output=True, **kwargs) # nosec
36
+ logger.info("--- BEGIN PROCESS STDOUT ---")
37
+ for line in result.stdout.decode("utf-8").strip().split("\n"):
38
+ logger.info(line)
39
+ logger.info("--- END PROCESS STDOUT ---")
40
+ logger.info("--- BEGIN PROCESS STDERR ---")
41
+ for line in result.stderr.decode("utf-8").strip().split("\n"):
42
+ logger.error(line)
43
+ logger.info("--- END PROCESS STDERR ---")
44
+ return result
45
+
46
+
47
+ def log_command_output(logger: logging.Logger, command: list[str], **kwargs) -> str:
48
+ logger.info(f"Run: {' '.join(command)}")
49
+ output = subprocess.check_output(command, **kwargs) # nosec
50
+ return output.decode("utf-8").strip()