py.lib

Yohn Y. 2022-07-17 Parent:1668cc57225b

27:54d7d7b9350b Go to Latest

py.lib/ssh_tools.py

. Исправление в форматировании

Download raw file
View source Diff to previous Annotate
History
1 # -*- coding: utf-8 -*-
2 import paramiko
3 from datetime import datetime
4
5 # --- OPTS --- #
6 TIMEOUT = 120
7
8 # --- CONST --- #
9 # Политика ключей узла
10 HOSTKEY_autoAdd = 0
11 HOSTKEY_warn = 1
12 HOSTKEY_reject = 2
13
14
15 def prompt_checker(user, host, buf):
16 signs = ['[', ']', '#', '<', '>', '$', '!', '%', '(', ')', '@', '+', ':', ';', ' ', '\t']
17 is_prompt = False
18 is_sep = True
19 end_buf = ''
20
21 ch_pos = 0
22 for ch_pos in range(1, len(buf) + 1):
23 if not is_prompt:
24 if buf[-ch_pos] == '\n':
25 end_buf = buf[-ch_pos] + end_buf
26 continue
27 is_prompt = True
28
29 if is_sep:
30 if buf[-ch_pos] in signs:
31 end_buf = buf[-ch_pos] + end_buf
32 continue
33 is_sep = False
34
35 if buf[-ch_pos] == '\n':
36 break
37
38 _startProc = len(buf) - ch_pos
39 _endProc = len(buf) - len(end_buf)
40 buf = buf[_startProc:_endProc]
41 search_cond = []
42
43 for ch_pos in [user, host]:
44 buf_l = ''
45 buf_r = ''
46 index = buf.find(ch_pos)
47 if index != -1:
48 for ch in buf[:index][::-1]:
49 if ch in signs:
50 buf_l = ch + buf_l
51 else:
52 break
53 index = index + len(ch_pos)
54 for ch in buf[index:]:
55 if ch in signs:
56 buf_r += ch
57 else:
58 break
59 search_cond.append(buf_l + ch_pos + buf_r)
60
61 def check(check_buf):
62 if check_buf.endswith(end_buf):
63 end_search = len(check_buf) - len(end_buf)
64 start_search = check_buf.rfind('\n', 0, len(check_buf) - len(end_buf))
65
66 if start_search == -1:
67 start_search = 0
68
69 _buf = check_buf[start_search:end_search]
70 for i in search_cond:
71 if _buf.find(i) == -1:
72 return False
73
74 return True
75
76 else:
77 return False
78
79 return check
80
81
82 class SSHError(Exception):
83 pass
84
85
86 class TOCheck(object):
87 """
88 Взято из: https://bitbucket.org/awgur/pylib/src/default/timetools.py
89 """
90 def __init__(self, timeout):
91 self.start = datetime.now()
92 self.timeout = timeout
93
94 def __call__(self, timeout=None):
95 if not timeout:
96 timeout = self.timeout
97
98 buf = datetime.now() - self.start
99
100 if buf.seconds > timeout:
101 return True
102
103 else:
104 return False
105
106 def __bool__(self):
107 return not self()
108
109
110 class SSHSubsh(object):
111 """Класс для облегчения работы с контекстами, когда приглашение коммандной оболочки
112 меняется в результате исполнения комманды.
113 """
114 def __init__(self, ssh, prompt):
115 self.ssh = ssh
116 self.prompt = SSHPromptCheck(prompt)
117
118 def __call__(self, cmd):
119 res = self.ssh(cmd, wait=False)
120 res += self.waitPrompt(self.prompt)
121 return res
122
123
124 class SSHPromptCheck(object):
125 """
126 Проверяет, появилось ли приглашение коммандной строки в выводе сервера
127
128 Принимает:
129 - prompt:
130 Строка, которая должна однозначно идентифицировать приглашение
131 коммандной строки, либо функция которой передаётся строка, и она
132 должна вернуть True, если строка завершается приглашением коммандной
133 строки либо False в обратном случае. Кроме того можно передат сам объект,
134 созданный заранее.
135 """
136 def __init__(self, prompt):
137 if isinstance(prompt, self.__class__):
138 self.prompt = prompt.prompt
139 self.check = prompt.check
140
141 else:
142 def is_func():
143 pass
144
145 if type(prompt) == type(is_func):
146 self.check = prompt
147 self.prompt = None
148
149 else:
150 self.prompt = prompt
151
152 if self.prompt is None:
153 raise SSHError('Prompt not set')
154
155 def check(self, buf):
156 if buf.endswith(self.prompt):
157 return True
158
159 else:
160 return False
161
162 def __call__(self, buf):
163 if isinstance(buf, SSHOutput):
164 return self.check(buf.out)
165
166 elif buf:
167 return self.check(buf)
168
169 else:
170 return False
171
172
173 class SSHOutput(object):
174 """Содержит вывод удалённого сервера и помогает управляться с ним.
175 """
176
177 def __init__(self, out='', err=''):
178 self.out = out
179 self.err = err
180
181 def error(self, buf):
182 """Добавить буфер к буферу сообщения об ошибке"""
183
184 self.err += buf
185
186 def __add__(self, buf):
187 if isinstance(buf, self.__class__):
188 self.out += buf.out
189 self.err += buf.err
190 else:
191 self.out += buf
192 return self
193
194 def __str__(self):
195 return self.out
196
197 def get_error(self):
198 return self.err
199
200 def __bool__(self):
201 if self.out or self.err:
202 return True
203 else:
204 return False
205
206 def __repr__(self):
207 buf = ''
208 buf += '---- OUTPUT -----------'
209 buf += self.out
210 buf += '---- ERRORS -----------'
211 buf += self.err
212 return buf
213
214 def __getitem__(self, key):
215 return self.out.splittines()[key]
216
217 def __iter__(self):
218 for line in self.out.splitlines():
219 yield line
220
221 def is_eol(self):
222 if self.out[-1] == '\n':
223 return True
224 else:
225 return False
226
227
228 class SSHCommunicate(object):
229 """\
230 Взаимодействие с командой ssh. В сосотоянии отправлять ей данные и принимать данные от неё.
231
232 Обеспечивает работу с exec_command в SSH
233 """
234
235 def __init__(self, stdin, stdout, stderr):
236 self.inC = stdin
237 self.out = stdout
238 self.err = stderr
239
240 def __call__(self, buf=None):
241 if buf is not None:
242 self.inC.write(buf)
243 return SSHOutput(self.out.read(), self.err.read())
244
245 def __str__(self):
246 return str(self())
247
248 def __repr__(self):
249 return repr(self())
250
251
252 class _SSH(object):
253 """\
254 Метакласс для SSH соединений
255
256 Принимает:
257 - host: Имя узла
258 - user: Имя пользователя
259 - passwd(не обязательно, если есть key): пароль пользователя
260 - key(не обязательно если есть passwd): файл ключа
261 - hostKeyPol[HOSTKEY_autoAdd | HOSTKEY_warn | HOSTKEY_reject]:
262 HOSTKEY_autoAdd(По умолчанию):
263 если ключ удалённого узла нам неизвестен, добавить, соединение разрешить.
264 HOSTKEY_warn: если ключ удалённого узла нам неизвестен, предупредить.
265 HOSTKEY_reject: если ключ удалённого узла нам неизвестен, разорвать соединение.
266 - timeout: таймаут операций с сервером.
267 """
268
269 def __init__(self, host, user, passwd=None, key=None,
270 host_key_pol=HOSTKEY_autoAdd, timeout=None
271 ):
272
273 if passwd is None and key is None:
274 raise SSHError('Auth type not set')
275
276 if host_key_pol == HOSTKEY_autoAdd:
277 _hostKeyPol = paramiko.AutoAddPolicy()
278 elif host_key_pol == HOSTKEY_warn:
279 _hostKeyPol = paramiko.WarningPolicy()
280 elif host_key_pol == HOSTKEY_reject:
281 _hostKeyPol = paramiko.RejectPolicy()
282 else:
283 raise SSHError('Unknown policy type')
284
285 self.ssh = paramiko.SSHClient()
286 self.ssh.set_missing_host_key_policy(_hostKeyPol)
287 if passwd:
288 self.ssh.connect(hostname=host, username=user, password=passwd, look_for_keys=False, allow_agent=False)
289 else:
290 self.ssh.connect(hostname=host, username=user, key_filename=key, look_for_keys=False, allow_agent=False)
291
292 if timeout is not None:
293 self.timeout = timeout
294
295 else:
296 self.timeout = TIMEOUT
297
298 self.user = user
299 self.host = host
300
301 def __del__(self):
302 try:
303 self.ssh.close()
304
305 except:
306 pass
307
308 del self.ssh
309
310
311 class InterSSH(_SSH):
312 """Класс Взаимодействия с сервером SSH
313
314 Кроме аргументов необходимых для метакласса принимает:
315 - prompt:
316 Параметр конструктора SSHPromptCheck. Принимается строго по имени.
317 """
318 def __init__(self, *a, **kwa):
319 prompt = None
320 if 'prompt' in kwa:
321 prompt = kwa['prompt']
322 del kwa['prompt']
323
324 _SSH.__init__(self, *a, **kwa)
325
326 self.shell = self.ssh.invoke_shell()
327 self.shell.settimeout(self.timeout)
328 if prompt is None:
329 prompt = self._get_prompt()
330 self.ready = True
331
332 else:
333 self.ready = False
334
335 self.prompt = SSHPromptCheck(prompt)
336
337 def __bool__(self):
338 """Готово ли соединение принимать новые команды"""
339 return self.ready
340
341 def _get_prompt(self):
342 """Возвращает объект проверки на приглашение комммандной строки из текщего
343 буфера вывода.
344 """
345 buf = ''
346 while self.shell.recv_ready():
347 buf += self.shell.recv(1024)
348
349 func = prompt_checker(self.user, self.host, buf)
350 return SSHPromptCheck(func)
351
352 def wait_prompt(self, prompt=None):
353 """\
354 Ожидание приглашения командной строки
355
356 Принимает необязательный аргумент: Экземпляр класса SSHPromptCheck, с помощью котороого
357 проверяется присутствие приглашения командного интерпретатора в выводе.
358 """
359
360 if prompt is None:
361 prompt = self.prompt
362
363 else:
364 prompt = SSHPromptCheck(prompt)
365
366 res = SSHOutput()
367 if not self.ready:
368 time_check = TOCheck(self.timeout)
369 while True:
370 while self.shell.recv_stderr_ready():
371 res.error(self.shell.recv_stderr(1024))
372
373 while self.shell.recv_ready():
374 res += self.shell.recv(1024)
375
376 if prompt(res):
377 break
378
379 if time_check():
380 raise SSHError('Command call timeout\n' + repr(res))
381
382 self.ready = True
383 return res
384
385 def __call__(self, cmd, wait=True):
386 """Запуск команды
387
388 Принимает:
389 - cmd: Комманда
390 - wait(необязательный аргумент): Признак ожидания завершения комманды.
391
392 Возвращает: Экземпляр класса SSHOutput
393 """
394 res = SSHOutput()
395 if not self.ready:
396 res += self.wait_prompt()
397 bytes = self.shell.send('%s\n' % cmd)
398 if bytes == 0:
399 raise SSHError('Channel closed')
400 if wait:
401 res += self.wait_prompt()
402 else:
403 self.ready = False
404 return res
405
406 def sub_shell(self, cmd, prompt=None):
407 """Возвращает объект SSHSubsh, через который удобно работать
408 интерпретаторами, запускаемыми из текущего, либо с ситуацией изменения приглашения
409 командного интерпретатора.
410 """
411 self(cmd, wait=False)
412 if prompt is None:
413 prompt = self._getPropt()
414 return SSHSubsh(self, prompt)
415
416 def close(self):
417 res = SSHOutput
418 if not self:
419 res += self.wait_prompt()
420 self.ssh.close()
421 return res
422
423
424 class SSH(_SSH):
425 """Класс призванный управлять выполнением единичных комманд, без привлечения контекста.
426
427 Принимает аргументы необходимые для метакласса _SSH
428 """
429 def __init__(self, *a, **kwa):
430 _SSH.__init__(self, *a, **kwa)
431
432 def __call__(self, cmd):
433 return SSHCommunicate(*self.ssh.exec_command(cmd, timeout=self.timeout))