Azalea
135 lines
3.8 KiB
Python
Executable File
135 lines
3.8 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
import importlib
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
import textwrap
|
|
from urllib.request import Request, urlopen
|
|
|
|
from color_utils import printc
|
|
|
|
|
|
def import_or_install(module: str, package: str | None = None):
|
|
try:
|
|
return importlib.import_module(module)
|
|
except ModuleNotFoundError:
|
|
subprocess.check_call([sys.executable, "-m", "pip", "install", package or module])
|
|
return importlib.import_module(module)
|
|
|
|
|
|
# Load github users
|
|
GITHUB_USERS = {'hykilpikonna', 'sauricat'}
|
|
ADDITIONAL_USERS_PATH = Path.home() / '.ssh' / 'authorized_github_users'
|
|
if ADDITIONAL_USERS_PATH.is_file():
|
|
GITHUB_USERS.update(set(ADDITIONAL_USERS_PATH.read_text('utf-8').split('\n')))
|
|
|
|
KEYS_PATH = Path.home() / '.ssh' / 'authorized_keys'
|
|
KEYS_PATH.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
def fetch_keys(user: str) -> set[str]:
|
|
req = Request(f"https://github.com/{user}.keys")
|
|
with urlopen(req) as response:
|
|
assert response.status == 200
|
|
resp = response.read().decode()
|
|
return {f"{l} {user}@github" for l in resp.strip().splitlines()}
|
|
|
|
|
|
def fetch_all_keys() -> set[str]:
|
|
return {l for u in GITHUB_USERS for l in fetch_keys(u)}
|
|
|
|
|
|
def normalize_key(key: str) -> str:
|
|
"""
|
|
Remove comment from keys
|
|
"""
|
|
return ' '.join(key.split(' ')[:2])
|
|
|
|
|
|
def update_ssh_keys():
|
|
all_keys = list(fetch_all_keys())
|
|
normalized_keys = {normalize_key(k) for k in all_keys}
|
|
|
|
existing_keys = set(KEYS_PATH.read_text('utf-8').strip().splitlines()) \
|
|
if KEYS_PATH.is_file() else set()
|
|
for k in existing_keys:
|
|
if normalize_key(k) not in normalized_keys:
|
|
all_keys.append(k)
|
|
|
|
len_diff = len(all_keys) - len(existing_keys)
|
|
if len_diff > 0:
|
|
printc(f"&aSSH Keys Updated! Added {len_diff} keys")
|
|
|
|
KEYS_PATH.write_text('\n'.join(all_keys))
|
|
|
|
|
|
def list_users():
|
|
printc(f'&6> Now you have {len(GITHUB_USERS)} GitHub users in authorized_keys')
|
|
printc(f'&6> They are: {", ".join(GITHUB_USERS)}')
|
|
|
|
|
|
def remove_user_keys(user: str):
|
|
keys = set(fetch_keys(user))
|
|
normalized_keys = {normalize_key(k) for k in keys}
|
|
|
|
existing_keys = set(KEYS_PATH.read_text('utf-8').strip().splitlines())
|
|
updated_keys = set(existing_keys)
|
|
for k in existing_keys:
|
|
if f'{user}@github' in k or normalize_key(k) in normalized_keys:
|
|
updated_keys.remove(k)
|
|
|
|
len_diff = len(existing_keys) - len(updated_keys)
|
|
if len_diff > 0:
|
|
printc(f"&aSSH Keys Updated! Removed {len_diff} keys")
|
|
|
|
KEYS_PATH.write_text('\n'.join(updated_keys))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.argv.pop(0)
|
|
|
|
# No args provided
|
|
if len(sys.argv) == 0:
|
|
update_ssh_keys()
|
|
exit(0)
|
|
|
|
# Add users
|
|
elif sys.argv[0] == 'add':
|
|
if len(sys.argv) == 1:
|
|
printc('&cUsage: update-ssh-keys.py add {github usernames...}')
|
|
exit(4)
|
|
usernames = sys.argv[1:]
|
|
|
|
GITHUB_USERS.update(set(usernames))
|
|
ADDITIONAL_USERS_PATH.write_text('\n'.join(GITHUB_USERS))
|
|
printc(f'&aAdded {", ".join(usernames)}!')
|
|
list_users()
|
|
update_ssh_keys()
|
|
|
|
# Remove users
|
|
elif sys.argv[0] == 'remove':
|
|
if len(sys.argv) == 1:
|
|
printc('&cUsage: update-ssh-keys.py remove {github usernames...}')
|
|
exit(4)
|
|
usernames = sys.argv[1:]
|
|
|
|
for u in usernames:
|
|
if u in GITHUB_USERS:
|
|
GITHUB_USERS.remove(u)
|
|
remove_user_keys(u)
|
|
ADDITIONAL_USERS_PATH.write_text('\n'.join(GITHUB_USERS))
|
|
printc(f'&aRemoved {", ".join(usernames)}!')
|
|
list_users()
|
|
|
|
# List users
|
|
elif sys.argv[0] == 'list':
|
|
list_users()
|
|
|
|
# Unknown argument
|
|
else:
|
|
print(textwrap.dedent("""
|
|
Usage: update-ssh-keys.py
|
|
- add/remove [github username]
|
|
- list
|
|
""")) |