-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.py
160 lines (123 loc) · 4.79 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
#!/usr/bin/env python3
import os
import json
import click
import subprocess
CONFIG_DIR = os.path.expanduser("~/.rocket")
CONFIG_FILE = os.path.join(CONFIG_DIR, "config.json")
def print_ascii_art():
art = r"""
_____ _ _
| __ \ | | | |
| |__) |___ ___| | _____| |_
| _ // _ \ / __| |/ / _ \ __|
| | \ \ (_) | (__| < __/ |_
|_| \_\___/ \___|_|\_\___|\__|
By Joeri Abbo
"""
print(art)
@click.group()
def rocket():
pass
@click.command()
@click.option('--username', default=None, help='Username for the connection')
@click.option('--host', default=None, help='Host for the connection')
@click.option('--nickname', default=None, help='Optional nickname for the connection')
@click.option('--through-proxy', is_flag=True, help='If set, the connection will go through the proxy')
def add(username, host, nickname, through_proxy):
print_ascii_art()
config = read_config()
if not username:
if 'default_username' in config:
use_default = click.confirm(f"Do you want to use the default username {config['default_username']}?")
if use_default:
username = config['default_username']
else:
username = click.prompt('Please enter a username')
else:
username = click.prompt('Please enter a username')
if not host:
host = click.prompt('Please enter a host')
connection = {'username': username, 'host': host, 'nickname': nickname or host}
if through_proxy:
if 'proxy' in config:
connection['through_proxy'] = True
else:
click.echo('No proxy found in the config. Please add a proxy using the add_proxy command.')
config['connections'] = config.get('connections', [])
config['connections'].append(connection)
write_config(config)
click.echo(f'Added connection: {connection}')
@click.command()
@click.option('--username', default=None, help='Username for the proxy')
@click.option('--host', default=None, help='Host for the proxy')
@click.option('--nickname', default=None, help='Optional nickname for the proxy')
def add_proxy(username, host, nickname):
print_ascii_art()
config = read_config()
if not username:
username = click.prompt('Please enter a username for the proxy')
if not host:
host = click.prompt('Please enter a host for the proxy')
proxy = {'username': username, 'host': host, 'nickname': nickname or host}
config['proxy'] = proxy
write_config(config)
click.echo(f'Added proxy: {proxy}')
@click.command()
def delete():
print_ascii_art()
config = read_config()
config['connections'] = []
write_config(config)
click.echo('All connections deleted')
@click.command()
@click.argument('nickname', required=False)
def launch(nickname):
print_ascii_art()
config = read_config()
if not config.get('connections'):
click.echo('No connections found, please add a connection first')
return
if nickname:
connection = next((c for c in config['connections'] if c.get('nickname') == nickname), None)
else:
connection_nicknames = [c.get('nickname') or c['host'] for c in config['connections']]
connection_nickname = click.prompt(
f"Please choose a connection by its nickname ({', '.join(connection_nicknames)})", type=str)
connection = next((c for c in config['connections'] if
c.get('nickname') == connection_nickname or c['host'] == connection_nickname), None)
if connection:
if 'proxy' in connection:
proxy = config['proxy']
command = f"ssh -J {proxy['username']}@{proxy['host']} {connection['username']}@{connection['host']}"
else:
command = f"ssh {connection['username']}@{connection['host']}"
click.echo(f'Launching: {command}')
subprocess.Popen([command], shell=True, executable='/bin/bash')
else:
click.echo(f"No connection found for nickname {nickname}")
@click.command()
def install():
print_ascii_art()
config = read_config()
default_username = click.prompt('Please enter the default username')
config['default_username'] = default_username
write_config(config)
click.echo(f'Default username set to: {default_username}')
def read_config():
if os.path.exists(CONFIG_FILE):
with open(CONFIG_FILE, 'r') as f:
return json.load(f)
return {}
def write_config(config):
if not os.path.exists(CONFIG_DIR):
os.makedirs(CONFIG_DIR)
with open(CONFIG_FILE, 'w') as f:
json.dump(config, f)
rocket.add_command(add)
rocket.add_command(add_proxy)
rocket.add_command(delete)
rocket.add_command(launch)
rocket.add_command(install)
if __name__ == '__main__':
rocket()