blob: f1786f8541eff83947d62925c41a15ddb91acee5 [file] [log] [blame]
koder aka kdanilov73084622016-11-16 21:51:08 +02001import time
2import errno
3import socket
4import logging
5import os.path
6import selectors
7from io import BytesIO
8from typing import cast, Dict, List, Set
9
10import paramiko
11
12from . import utils
13from .ssh_utils import ConnCreds, IPAddr
14
15logger = logging.getLogger("wally")
16NODE_KEYS = {} # type: Dict[IPAddr, paramiko.RSAKey]
17
18
19def set_key_for_node(host_port: IPAddr, key: bytes) -> None:
20 with BytesIO(key) as sio:
21 NODE_KEYS[host_port] = paramiko.RSAKey.from_private_key(sio)
22
23
24def connect(creds: ConnCreds,
25 conn_timeout: int = 60,
26 tcp_timeout: int = 15,
27 default_banner_timeout: int = 30) -> paramiko.SSHClient:
28
29 ssh = paramiko.SSHClient()
30 ssh.load_host_keys('/dev/null')
31 ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
32 ssh.known_hosts = None
33
34 end_time = time.time() + conn_timeout # type: float
35
36 while True:
37 try:
38 time_left = end_time - time.time()
39 c_tcp_timeout = min(tcp_timeout, time_left)
40
41 banner_timeout_arg = {} # type: Dict[str, int]
42 if paramiko.__version_info__ >= (1, 15, 2):
43 banner_timeout_arg['banner_timeout'] = int(min(default_banner_timeout, time_left))
44
45 if creds.passwd is not None:
46 ssh.connect(creds.addr.host,
47 timeout=c_tcp_timeout,
48 username=creds.user,
49 password=cast(str, creds.passwd),
50 port=creds.addr.port,
51 allow_agent=False,
52 look_for_keys=False,
53 **banner_timeout_arg)
54 elif creds.key_file is not None:
55 ssh.connect(creds.addr.host,
56 username=creds.user,
57 timeout=c_tcp_timeout,
58 key_filename=cast(str, creds.key_file),
59 look_for_keys=False,
60 port=creds.addr.port,
61 **banner_timeout_arg)
62 elif creds.key is not None:
63 with BytesIO(creds.key) as sio:
64 ssh.connect(creds.addr.host,
65 username=creds.user,
66 timeout=c_tcp_timeout,
67 pkey=paramiko.RSAKey.from_private_key(sio),
68 look_for_keys=False,
69 port=creds.addr.port,
70 **banner_timeout_arg)
71 elif (creds.addr.host, creds.addr.port) in NODE_KEYS:
72 ssh.connect(creds.addr.host,
73 username=creds.user,
74 timeout=c_tcp_timeout,
75 pkey=NODE_KEYS[creds.addr],
76 look_for_keys=False,
77 port=creds.addr.port,
78 **banner_timeout_arg)
79 else:
80 key_file = os.path.expanduser('~/.ssh/id_rsa')
81 ssh.connect(creds.addr.host,
82 username=creds.user,
83 timeout=c_tcp_timeout,
84 key_filename=key_file,
85 look_for_keys=False,
86 port=creds.addr.port,
87 **banner_timeout_arg)
88 return ssh
89 except paramiko.PasswordRequiredException:
90 raise
91 except (socket.error, paramiko.SSHException):
92 if time.time() > end_time:
93 raise
94 time.sleep(1)
95
96
97def wait_ssh_available(addrs: List[IPAddr],
98 timeout: int = 300,
99 tcp_timeout: float = 1.0) -> None:
100
101 addrs_set = set(addrs) # type: Set[IPAddr]
102
103 for _ in utils.Timeout(timeout):
104 selector = selectors.DefaultSelector() # type: selectors.BaseSelector
105 with selector:
106 for addr in addrs_set:
107 sock = socket.socket()
108 sock.setblocking(False)
109 try:
110 sock.connect(addr)
111 except BlockingIOError:
112 pass
113 selector.register(sock, selectors.EVENT_READ, data=addr)
114
115 etime = time.time() + tcp_timeout
116 ltime = etime - time.time()
117 while ltime > 0:
118 # convert to greater or equal integer
119 for key, _ in selector.select(timeout=int(ltime + 0.99999)):
120 selector.unregister(key.fileobj)
121 try:
122 key.fileobj.getpeername() # type: ignore
123 addrs_set.remove(key.data)
124 except OSError as exc:
125 if exc.errno == errno.ENOTCONN:
126 pass
127 ltime = etime - time.time()
128
129 if not addrs_set:
130 break
131
132