/
__main__.py
165 lines (145 loc) · 5.02 KB
/
__main__.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
import os
import sys
from threading import Event, Thread
from typing import Callable
import resotolib.proc
from resotolib.args import ArgumentParser, Namespace
from resotolib.core import resotocore, add_args as core_add_args, resotocore_is_up
from resotolib.core.ca import TLSData
from resotolib.event import add_event_listener, Event as ResotoEvent, EventType
from resotolib.jwt import add_args as jwt_add_args
from resotolib.logger import log, setup_logger, add_args as logging_add_args
from resotoshell.promptsession import PromptSession
from resotoshell.shell import Shell
from rich.console import Console
from resotoclient import ResotoClient
def main() -> None:
resotolib.proc.parent_pid = os.getpid()
resotolib.proc.initializer()
setup_logger("resotoshell", json_format=False)
arg_parser = ArgumentParser(
description="resoto shell", env_args_prefix="RESOTOSHELL_"
)
core_add_args(arg_parser)
add_args(arg_parser)
logging_add_args(arg_parser)
jwt_add_args(arg_parser)
TLSData.add_args(arg_parser, ca_only=True)
args = arg_parser.parse_args()
if not resotocore_is_up(resotocore.http_uri):
log.fatal(f"resotocore is not online at {resotocore.http_uri}")
sys.exit(1)
client = ResotoClient(
url=resotocore.http_uri,
psk=args.psk,
custom_ca_cert_path=args.ca_cert,
verify=args.verify_certs,
)
client.start()
if args.stdin or not sys.stdin.isatty():
handle_from_stdin(client)
else:
repl(client, args)
client.shutdown()
resotolib.proc.kill_children(resotolib.proc.SIGTERM, ensure_death=True)
log.debug("Shutdown complete")
sys.exit(0)
def repl(
client: ResotoClient,
args: Namespace,
) -> None:
shutdown_event = Event()
shell = Shell(client, True, detect_color_system(args))
session = PromptSession()
log.debug("Starting interactive session")
def shutdown(event: ResotoEvent) -> None:
shutdown_event.set()
kt = Thread(target=resotolib.proc.delayed_exit, name="shutdown")
kt.start()
add_event_listener(EventType.SHUTDOWN, shutdown)
# send the welcome command to the core
shell.handle_command("welcome")
while not shutdown_event.is_set():
try:
command = session.prompt()
if command == "":
continue
if command == "quit":
shutdown_event.set()
continue
shell.handle_command(command)
except KeyboardInterrupt:
pass
except EOFError:
shutdown_event.set()
except (RuntimeError, ValueError) as e:
log.error(e)
except Exception:
log.exception("Caught unhandled exception while processing CLI command")
def handle_from_stdin(client: ResotoClient) -> None:
shell = Shell(client, False, "monochrome")
log.debug("Reading commands from STDIN")
try:
for command in sys.stdin.readlines():
command = command.rstrip()
shell.handle_command(command)
except KeyboardInterrupt:
pass
except (RuntimeError, ValueError) as e:
log.error(e)
except Exception:
log.exception("Caught unhandled exception while processing CLI command")
def detect_color_system(args: Namespace) -> str:
if args.no_color:
return "monochrome"
else:
lookup = {
None: "monochrome",
"standard": "standard",
"256": "eight_bit",
"truecolor": "truecolor",
"windows": "legacy_windows",
}
cs = lookup.get(Console().color_system, "standard")
log.debug(f"Detected color system is: {cs}")
return cs
def add_args(arg_parser: ArgumentParser) -> None:
def is_file(message: str) -> Callable[[str], str]:
def check_file(path: str) -> str:
if os.path.isfile(path):
return path
else:
raise AttributeError(f"{message}: path {path} is not a directory!")
return check_file
arg_parser.add_argument(
"--resotocore-section",
help="All queries are interpreted with this section name. If not set, the server default is used.",
dest="resotocore_section",
)
arg_parser.add_argument(
"--resotocore-graph",
help="The name of the graph to use by default. If not set, the server default is used.",
dest="resotocore_graph",
)
arg_parser.add_argument(
"--download-directory",
help="If files are received, they are written to this directory.",
default=".",
dest="download_directory",
)
arg_parser.add_argument(
"--no-color",
help="Output should be rendered plain without any color escape sequences.",
dest="no_color",
action="store_true",
default=False,
)
arg_parser.add_argument(
"--stdin",
help="Read from STDIN instead of opening a shell",
dest="stdin",
action="store_true",
default=False,
)
if __name__ == "__main__":
main()