Quickly change a ssh config HostName as well as reading and configuring the .ssh/config file.

ssh_tools.py

#!/usr/bin/env python

"""
Quickly change a ssh config HostName as well as reading and configuring the .ssh/config file.

Also supports syncing a pycharm http client env file.
"""

import fire
import shlex
import time
import os
import json
import shutil
import tempfile
import contextlib

try:
    import pyperclip as pc
except ImportError:
    try:
        import pyperclip3 as pc
    except ImportError:
        pc = None


CONFIGNAME = os.path.expanduser(os.environ.get("SSH_CONFIG", "~/.ssh/config"))
PYCHARM_ENV = os.path.expanduser(
    os.environ.get("PYCHARM_HTTP_ENV", "~/scripts/Requests/http-client.env.json")
)
YES = ["y", "yes", "true", "1", ""]


def create_config(
    name: str = None,
    hostname: str = None,
    forward_agent: bool = None,
    user: str = None,
    configname: str = CONFIGNAME,
):
    if name is None:
        name = input(f"Enter config name: ").strip()
    if not name:
        return
    if hostname is None:
        hostname = input("HostName: ").strip()
        if "\n" in hostname:
            hostname = ""
    if not hostname:
        return

    if forward_agent is None:
        forward_agent = input("ForwardAgent: ").strip().lower() in YES
    forward_agent = "yes" if forward_agent else "no"
    if user is None:
        user = input("User: ").strip()

    with atomic_open(configname, "a") as writer:
        writer.write(
            "\n"
            f"Host {name}\n"
            f"    HostName {hostname}\n"
            f"    ForwardAgent {forward_agent}\n"
            f"    User {user}\n"
        )


def show_config(name: str = None, configname: str = CONFIGNAME):
    with open(configname, "r") as reader:
        if name is None:
            config = reader.read()
        else:
            s, e, lines = get_config(name, configname=configname)
            config = "".join(lines[s:e])
    print(config)


def get_config(name: str = None, configname: str = CONFIGNAME):
    if name is None:
        name = input(f"Enter config name: ").strip()

    with open(configname, "r") as reader:
        lines = reader.readlines()

    i = -1
    end = -1
    for j, line in enumerate(lines):
        if line.strip() == f"Host {name}":
            i = j
        elif i > 0 and line.strip() == "":
            end = j
            break

    return i, end, lines


def print_config(name: str = None, configname: str = CONFIGNAME):
    if name is None:
        name = input(f"Enter config name: ").strip()

    s, e, lines = get_config(name, configname=configname)
    config = lines[s:e]
    if config:
        print("".join(config))
    else:
        print(f'"{name}" config not found!')


def set_config(name: str = None, hostname: str = None, configname: str = CONFIGNAME):
    if name is None:
        name = input(f"Enter config name: ").strip()
    if not name:
        return

    s, e, lines = get_config(name, configname=configname)
    if s < 0:
        cmd = (
            input(f"'Host {name}' not found!\nWould you like to create (Y/n)? ").strip()
            in YES
        )
        if cmd:
            create_config(name, configname=configname)
        return

    # Modify
    if hostname is None:
        hostname = input("HostName: ").strip()
        if "\n" in hostname:
            hostname = ""
    if not hostname:
        return
    for j, line in enumerate(lines[s:]):
        if "HostName" in line:
            lines[s + j] = f"    HostName {hostname}\n"
            break

    with atomic_open(configname, "w") as writer:
        writer.writelines(lines)

    # Load into Pycharm enf
    if os.path.exists(PYCHARM_ENV):
        load_config(name, configname=configname)


def load_config(name: str = None, configname: str = CONFIGNAME):
    """Load ssh into pycharm http requests file"""
    if name is None:
        name = input(f"Enter config name: ").strip()
    if not name:
        return

    hostname = copy_config(name, configname=configname)

    with open(PYCHARM_ENV, "r") as reader:
        data = json.load(reader)

    data[name]["host"] = hostname

    with open(PYCHARM_ENV, "w") as writer:
        json.dump(data, writer, indent=2)


def copy_config(
    name: str = None, full_config: bool = False, configname: str = CONFIGNAME
):
    if name is None:
        name = input(f"Enter config name: ").strip()
    if not name:
        return

    txt = None
    found = False
    with open(configname, "r") as reader:
        for line in reader:
            if found and full_config:
                if "Host " in line:
                    break
                txt += line
            elif found and "HostName" in line:
                txt = line.strip().split("HostName")[-1].strip()
                break
            if line.strip() == f"Host {name}":
                found = True
                if full_config:
                    txt = line
        else:
            raise ValueError(f"'Host {name}' not found!")

    if pc is not None:
        pc.copy(txt)
    return txt


@contextlib.contextmanager
def atomic_open(filename, mode="r", suffix="", dir=None, fsync=True):
    """Context for temporary file.

    Will find a free temporary filename upon entering
    and will try to delete the file on leaving, even in case of an exception.

    Parameters
    ----------
    suffix : string
        optional file suffix
    dir : string
        optional directory to save temporary file in
    """
    tf = tempfile.NamedTemporaryFile(mode, delete=False, suffix=suffix, dir=dir)
    try:
        tf.file.close()
        with open(tf.name, mode) as f:
            yield f
            if fsync:
                f.flush()
                os.fsync(f.fileno())
        shutil.copy(tf.name, filename)
    finally:
        try:
            os.remove(tf.name)
        except OSError as e:
            if e.errno == 2:
                pass
            else:
                raise


def interactive():
    while True:
        try:
            cmd = input("Enter command: ").strip()
            if cmd == "quit" or cmd == "exit":
                break

            cmd = shlex.split(cmd)
            sys.argv = sys.argv[0:1] + cmd

            fire.Fire(CMDS)
        except KeyboardInterrupt:
            print("User exited the application")
            break
        except (ValueError, TypeError, SystemExit, Exception) as err:
            print(err)
            time.sleep(0.3)


CMDS = {
    "get": print_config,
    "set": set_config,
    "load": load_config,
    "create": create_config,
    "show": show_config,
    "interactive": interactive,
}
if pc is not None:
    CMDS["copy"] = copy_config


if __name__ == "__main__":
    import sys

    if len(sys.argv) == 1:
        interactive()
    else:
        fire.Fire(CMDS)