koder aka kdanilov | 7308462 | 2016-11-16 21:51:08 +0200 | [diff] [blame^] | 1 | import time |
| 2 | import errno |
| 3 | import socket |
| 4 | import logging |
| 5 | import os.path |
| 6 | import selectors |
| 7 | from io import BytesIO |
| 8 | from typing import cast, Dict, List, Set |
| 9 | |
| 10 | import paramiko |
| 11 | |
| 12 | from . import utils |
| 13 | from .ssh_utils import ConnCreds, IPAddr |
| 14 | |
| 15 | logger = logging.getLogger("wally") |
| 16 | NODE_KEYS = {} # type: Dict[IPAddr, paramiko.RSAKey] |
| 17 | |
| 18 | |
| 19 | def 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 | |
| 24 | def 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 | |
| 97 | def 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 | |