Skip to content

Commit 349f8bb

Browse files
jessebotjessebot
committed
add pretty help menu and pretty logging
1 parent 73f1acc commit 349f8bb

File tree

2 files changed

+127
-91
lines changed

2 files changed

+127
-91
lines changed

‎awth/__init__.py‎

Lines changed: 118 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from botocore.exceptions import ClientError, ParamValidationError
88

99
# import bw
10-
import click
10+
from click import option, command, Choice
1111
try:
1212
import configparser
1313
from configparser import NoOptionError, NoSectionError
@@ -21,106 +21,139 @@
2121
from os import path, environ, makedirs
2222
from pathlib import Path
2323
from rich.prompt import Confirm, Prompt
24+
from rich.logging import RichHandler
25+
from .help_text import RichCommand, options_help
2426
import sys
2527

26-
logger = logging.getLogger('awth')
2728
AWS_CREDS_PATH = path.expanduser(environ.get('AWS_SHARED_CREDENTIALS_FILE', '~/.aws/credentials'))
29+
HELP = options_help()
30+
LOG_LEVEL = 'warn'
31+
LOG_FILE = None
2832

2933

30-
def setup_logger(level=logging.DEBUG):
34+
def setup_logger(level="", log_file=""):
3135
"""
32-
set up basic logger
36+
Sets up rich logger and stores the values for it in a db for future import
37+
in other files. Returns logging.getLogger("rich")
3338
"""
34-
stdout_handler = logging.StreamHandler(stream=sys.stdout)
35-
stdout_handler.setFormatter(
36-
logging.Formatter('%(levelname)s - %(message)s'))
37-
stdout_handler.setLevel(level)
38-
logger.addHandler(stdout_handler)
39-
logger.setLevel(level)
40-
41-
42-
@click.command()
43-
@click.option('--device',
44-
required=False,
45-
metavar='arn:aws:iam::123456788990:mfa/dudeman',
46-
help="The MFA Device ARN. This value can also be "
47-
"provided via the environment variable 'MFA_DEVICE' or"
48-
" the ~/.aws/credentials variable 'aws_mfa_device'.")
49-
@click.option('--duration',
50-
type=int,
51-
help="The duration, in seconds, that the temporary "
52-
"credentials should remain valid. Minimum value: "
53-
"900 (15 minutes). Maximum: 129600 (36 hours). "
54-
"Defaults to 43200 (12 hours), or 3600 (one "
55-
"hour) when using '--assume-role'. This value "
56-
"can also be provided via the environment "
57-
"variable 'MFA_STS_DURATION'. ")
58-
@click.option('--profile',
59-
help="If using profiles, specify the name here. The "
60-
"default profile name is 'default'. The value can "
61-
"also be provided via the environment variable "
62-
"'AWS_PROFILE'.",
63-
required=False)
64-
@click.option('--long-term-suffix', '--long-suffix',
65-
help="The suffix appended to the profile name to"
66-
"identify the long term credential section",
67-
required=False)
68-
@click.option('--short-term-suffix', '--short-suffix',
69-
help="The suffix appended to the profile name to"
70-
"identify the short term credential section",
71-
required=False)
72-
@click.option('--assume-role', '--assume',
73-
metavar='arn:aws:iam::123456788990:role/RoleName',
74-
help="The ARN of the AWS IAM Role you would like to "
75-
"assume, if specified. This value can also be provided"
76-
" via the environment variable 'MFA_ASSUME_ROLE'",
77-
required=False)
78-
@click.option('--role-session-name',
79-
help="Friendly session name required when using "
80-
"--assume-role",
81-
default=getpass.getuser(),
82-
required=False)
83-
@click.option('--force',
84-
help="Refresh credentials even if currently valid.",
85-
required=False)
86-
@click.option('--log-level',
87-
type=click.Choice(['CRITICAL', 'ERROR', 'WARNING', 'INFO', 'DEBUG', 'NOTSET'], case_sensitive=False),
88-
help="Set log level",
89-
required=False,
90-
default='DEBUG')
91-
@click.option('--setup',
92-
help="Setup a new log term credentials section",
93-
is_flag=bool,
94-
required=False)
95-
@click.option('--token',
96-
help="Provide MFA token as an argument",
97-
required=False,
98-
default=None)
99-
@click.option('--region',
100-
help="AWS STS Region",
101-
required=False,
102-
type=str)
103-
@click.option('--keychain',
104-
is_flag=bool,
105-
help="Use system keychain to store or retrieve long term credentials",
106-
required=False)
39+
# determine logging level
40+
if not level:
41+
level = LOG_LEVEL
42+
43+
log_level = getattr(logging, level.upper(), None)
44+
45+
# these are params to be passed into logging.basicConfig
46+
opts = {'level': log_level, 'format': "%(message)s", 'datefmt': "[%X]"}
47+
48+
# we only log to a file if one was passed into config.yml or the cli
49+
if not log_file:
50+
log_file = LOG_FILE
51+
52+
# rich typically handles much of this but we don't use rich with files
53+
if log_file:
54+
opts['filename'] = log_file
55+
opts['format'] = "%(asctime)s %(levelname)s %(funcName)s: %(message)s"
56+
else:
57+
rich_handler_opts = {'rich_tracebacks': True}
58+
# 10 is the DEBUG logging level int value
59+
if log_level == 10:
60+
# log the name of the function if we're in debug mode :)
61+
opts['format'] = "[bold]%(funcName)s()[/bold]: %(message)s"
62+
rich_handler_opts['markup'] = True
63+
64+
opts['handlers'] = [RichHandler(**rich_handler_opts)]
65+
66+
# this uses the opts dictionary as parameters to logging.basicConfig()
67+
logging.basicConfig(**opts)
68+
69+
if log_file:
70+
return None
71+
else:
72+
return logging.getLogger("rich")
73+
74+
75+
@command(cls=RichCommand)
76+
@option('--device',
77+
required=False,
78+
metavar='arn:aws:iam::123456788990:mfa/dudeman',
79+
help="The MFA Device ARN. This value can also be "
80+
"provided via the environment variable 'MFA_DEVICE' or"
81+
" the ~/.aws/credentials variable 'aws_mfa_device'.")
82+
@option('--duration',
83+
type=int,
84+
help="The duration, in seconds, that the temporary "
85+
"credentials should remain valid. Minimum value: "
86+
"900 (15 minutes). Maximum: 129600 (36 hours). "
87+
"Defaults to 43200 (12 hours), or 3600 (one "
88+
"hour) when using '--assume-role'. This value "
89+
"can also be provided via the environment "
90+
"variable 'MFA_STS_DURATION'. ")
91+
@option('--profile',
92+
help="If using profiles, specify the name here. The "
93+
"default profile name is 'default'. The value can "
94+
"also be provided via the environment variable "
95+
"'AWS_PROFILE'.",
96+
required=False)
97+
@option('--long-term-suffix', '--long-suffix',
98+
help="The suffix appended to the profile name to"
99+
"identify the long term credential section",
100+
required=False)
101+
@option('--short-term-suffix', '--short-suffix',
102+
help="The suffix appended to the profile name to"
103+
"identify the short term credential section",
104+
required=False)
105+
@option('--assume-role', '--assume',
106+
metavar='arn:aws:iam::123456788990:role/RoleName',
107+
help="The ARN of the AWS IAM Role you would like to "
108+
"assume, if specified. This value can also be provided"
109+
" via the environment variable 'MFA_ASSUME_ROLE'",
110+
required=False)
111+
@option('--role-session-name',
112+
help="Friendly session name required when using "
113+
"--assume-role",
114+
default=getpass.getuser(),
115+
required=False)
116+
@option('--force',
117+
help="Refresh credentials even if currently valid.",
118+
required=False)
119+
@option('--log_level',
120+
type=Choice(['CRITICAL', 'ERROR', 'WARNING', 'INFO', 'DEBUG', 'NOTSET'],
121+
case_sensitive=False),
122+
help="Set log level",
123+
required=False,
124+
default='DEBUG')
125+
@option('--setup',
126+
help="Setup a new log term credentials section",
127+
is_flag=bool,
128+
required=False)
129+
@option('--token',
130+
help="Provide MFA token as an argument",
131+
required=False,
132+
default=None)
133+
@option('--region',
134+
help="AWS STS Region",
135+
required=False,
136+
type=str)
137+
@option('--keychain',
138+
is_flag=bool,
139+
help="Use system keychain to store or retrieve long term credentials",
140+
required=False)
107141
def main(device: str,
108142
duration: int,
109143
profile: str,
110144
long_term_suffix: str,
111145
short_term_suffix: str,
112146
assume_role: str,
113147
role_session_name: str,
114-
force: bool,
115-
log_level: str,
116-
setup: bool,
148+
force: bool = False,
149+
log_level: str = "INFO",
150+
setup: bool = False,
117151
token: str = "",
118152
region: str = "eu-central-1",
119153
keychain: bool = False):
120154

121155
# set up logging before we begin
122-
level = getattr(logging, log_level)
123-
setup_logger(level)
156+
logger = setup_logger(log_level)
124157

125158
if not path.isfile(AWS_CREDS_PATH):
126159
create_credentials_file = Confirm.ask(
@@ -145,6 +178,7 @@ def main(device: str,
145178
return
146179

147180
validate(config,
181+
logger,
148182
profile,
149183
long_term_suffix,
150184
short_term_suffix,
@@ -155,7 +189,7 @@ def main(device: str,
155189
force)
156190

157191

158-
def get_config(aws_creds_path: str = ""):
192+
def get_config(logger, aws_creds_path: str = ""):
159193
"""
160194
get the configuration and parse it
161195
"""
@@ -172,6 +206,7 @@ def get_config(aws_creds_path: str = ""):
172206

173207

174208
def validate(config,
209+
logger,
175210
profile: str = "",
176211
long_term_suffix: str = "",
177212
short_term_suffix: str = "",

‎awth/help_text.py‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@
22
help text functions for the onboardme cli
33
"""
44
# file for rich printing
5+
import awth
56
import click
7+
from importlib.metadata import version as get_version
68
from rich.console import Console
79
from rich.highlighter import RegexHighlighter
810
from rich.panel import Panel
@@ -11,8 +13,7 @@
1113
from rich.theme import Theme
1214

1315
# custom local module
14-
from .constants import (VERSION)
15-
16+
VERSION = get_version("awth")
1617

1718
# this is for creating new help text svgs for the READMEs
1819
RECORD = False
@@ -89,20 +90,20 @@ class OptionHighlighter(RegexHighlighter):
8990

9091
highlighter = OptionHighlighter()
9192

92-
console = Console(theme=Theme({"option": "cornflower_blue",
93-
"switch": "deep_sky_blue1",
94-
"meta": "light_steel_blue",
93+
console = Console(theme=Theme({"option": "green",
94+
"switch": "magenta",
95+
"meta": "blue",
9596
"unstable": "italic cyan"}),
9697
highlighter=highlighter, record=RECORD)
9798

98-
title = "☁️ [cornflower_blue]awth[/] 🗝️\n"
99+
title = "🌤️ [green]awth[/] 🗝️\n"
99100
desc = (
100101
"[steel_blue]Authenticate to AWS using MFA.")
101102

102103
console.print(title + desc, justify="center")
103104

104-
console.print("\n[b]Usage[/]: [royal_blue1]awth[/] " +
105-
"[cornflower_blue][OPTIONS]\n")
105+
console.print("\n[b]Usage[/]: [magenta]awth[/] " +
106+
"[green][OPTIONS]\n")
106107

107108
options_table = Table(highlight=True, box=None, show_header=False,
108109
row_styles=["", "dim"],

0 commit comments

Comments
 (0)