ssh in python

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
from abc import ABCMeta, abstractmethod
from paramiko import SSHClient, AutoAddPolicy
from subprocess import run


class ConnectorFactory:
    def __init__(self, config):
        self.config = config
        if 'local' in config:
            self.local = config['local']
        else:
            self.local = False

    def get(self):
        if self.local:
            return LocalConnector(self.config)
        else:
            return ServerConnector(self.config)
class Connector(metaclass=ABCMeta):
    def __init__(self, config):
        self.config    = config
        self.id        = f" {config['username']}@{config['hostname']}:{config['port']} "
        self.bg_format = "nohup %s >> nohup.log 2>&1 &"

    def cmds(self, cmds):
        cmds = cmds.strip()
        proc_cmds = self._process_cmds(cmds)
        _, stdout, stderr = self._execute_cmds(proc_cmds)
        self._process_result(stdout, stderr)
    @abstractmethod
    def _process_cmds(self, cmds):
        pass
    @abstractmethod
    def _execute_cmds(self, cmds):
        pass
    @abstractmethod
    def _process_result(self, stdout, stderr):
        pass
class ServerConnector(Connector):
    def __init__(self, config):
        super().__init__(config)
        self.client = SSHClient()
        self.client.set_missing_host_key_policy(AutoAddPolicy)

    def cmds(self, cmds):
        self.client.connect(**self.config)
        super().cmds(cmds)
        self.client.close()

    def _process_cmds(self, cmds):
        return '\n'.join([self.bg_format % c.replace('&', '') if '&' in c else c for c in cmds.split('\n')])

    def _execute_cmds(self, cmds):
        return self.client.exec_command(cmds)

    def _process_result(self, stdout, stderr):
        for fd in (stdout, stderr):
            if lines := fd.readlines():
                print()
                print(f"┌{self.id:─<80}┐")
                for line in lines:
                    print('│', line.strip())
                print(f"└{'─'*80}┘")
class LocalConnector(Connector):
    def _process_cmds(self, cmds):
        return ';'.join([self.bg_format % c.replace('&', '') if '&' in c else c for c in cmds.split('\n')])

    def _execute_cmds(self, cmds):
        stds = run(cmds, shell=True, executable='/bin/bash', capture_output=True)
        return None, stds.stdout, stds.stderr

    def _process_result(self, stdout, stderr):
        for fd in (stdout, stderr):
            if lines := fd.decode().strip():
                print()
                print(f"┌{self.id:─<80}┐")
                for line in lines.split('\n'):
                    print('│', line)
                print(f"└{'─'*80}┘")


if __name__ == '__main__':
    cmds = """
    source /opt/conda/bin/activate rapids
    which python
    """

    configs = [
        dict(hostname='123.456.78.910', port=26030, username='root'),
        dict(hostname='123.456.78.911', port=10022, username='root', local=True)
    ]

    for config in configs:
        conn = ConnectorFactory(config).get()
        conn.cmds(cmds)