/
cmd_ssh_keys.py
58 lines (45 loc) · 1.92 KB
/
cmd_ssh_keys.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
import json
import random
import string
from typing import List, Optional
from lightning_cloud.openapi import V1CreateSSHPublicKeyRequest, V1SSHPublicKey
from rich.console import Console
from rich.table import Table
from lightning_app.cli.core import Formatable
from lightning_app.utilities.network import LightningClient
class _SSHKeyList(Formatable):
def __init__(self, ssh_keys: List[V1SSHPublicKey]):
self.ssh_keys = ssh_keys
def as_json(self) -> str:
return json.dumps(self.ssh_keys)
def as_table(self) -> Table:
table = Table("id", "public_key", "created", show_header=True, header_style="bold green")
for ssh_key in self.ssh_keys:
table.add_row(
ssh_key.id,
ssh_key.public_key,
ssh_key.created_at.strftime("%Y-%m-%d"),
)
return table
class _SSHKeyManager:
"""_SSHKeyManager implements API calls specific to Lightning AI SSH-Keys."""
def __init__(self) -> None:
self.api_client = LightningClient()
def get_ssh_keys(self) -> _SSHKeyList:
resp = self.api_client.s_sh_public_key_service_list_ssh_public_keys()
return _SSHKeyList(resp.ssh_keys)
def list(self) -> None:
ssh_keys = self.get_ssh_keys()
console = Console()
console.print(ssh_keys.as_table())
def add_key(self, public_key: str, name: Optional[str], comment: Optional[str]) -> None:
key_name = name if name is not None else "-".join(random.choice(string.ascii_lowercase) for _ in range(5))
self.api_client.s_sh_public_key_service_create_ssh_public_key(
V1CreateSSHPublicKeyRequest(
name=key_name,
public_key=public_key,
comment=comment if comment is not None else key_name,
)
)
def remove_key(self, key_id: str) -> None:
self.api_client.s_sh_public_key_service_delete_ssh_public_key(key_id)