-
Notifications
You must be signed in to change notification settings - Fork 5.5k
/
Copy pathvt_helper.py
136 lines (115 loc) · 4.38 KB
/
vt_helper.py
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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
# -*- coding: utf-8 -*-
'''
salt.utils.vt_helper
~~~~~~~~~~~~~~~~~~~~
VT Helper
This module provides the SSHConnection to expose an SSH connection object
allowing users to programatically execute commands on a remote server using
Salt VT.
'''
from __future__ import absolute_import
# Import python libs
import logging
import os
import re
# Import salt's Libs
from .vt import Terminal, TerminalException
SSH_PASSWORD_PROMPT_RE = re.compile(r'(?:.*)[Pp]assword(?: for .*)?:', re.M)
KEY_VALID_RE = re.compile(r'.*\(yes\/no\).*')
log = logging.getLogger(__name__)
class SSHConnection(object):
'''
SSH Connection to a remote server.
'''
def __init__(self,
username='salt',
password='password',
host='localhost',
key_accept=False,
prompt=r'(Cmd)',
passwd_retries=3,
linesep=os.linesep):
'''
Establishes a connection to the remote server.
The format for parameters is:
username (string): The username to use for this
ssh connection. Defaults to root.
password (string): The password to use for this
ssh connection. Defaults to password.
host (string): The host to connect to.
Defaults to localhost.
key_accept (boolean): Should we accept this host's key
and add it to the known_hosts file? Defaults to False.
prompt (string): The shell prompt (regex) on the server.
Prompt is compiled into a regular expression.
Defaults to (Cmd)
passwd_retries (int): How many times should I try to send the password?
Defaults to 3.
linesep (string): The line separator to use when sending
commands to the server. Defaults to os.linesep.
'''
self.conn = Terminal(
'ssh {0}@{1}'.format(username, host),
shell=True,
log_stdout=True,
log_stdout_level='trace',
log_stderr=True,
log_stderr_level='trace',
stream_stdout=False,
stream_stderr=False)
sent_passwd = 0
self.prompt_re = re.compile(prompt)
self.linesep = linesep
while self.conn.has_unread_data:
stdout, stderr = self.conn.recv()
if stdout and SSH_PASSWORD_PROMPT_RE.search(stdout):
if not password:
log.error('Failure while authentication.')
raise TerminalException(
'Permission denied, no authentication information')
if sent_passwd < passwd_retries:
self.conn.sendline(password, self.linesep)
sent_passwd += 1
continue
else:
# asking for a password, and we can't seem to send it
raise TerminalException('Password authentication failed')
elif stdout and KEY_VALID_RE.search(stdout):
# Connecting to this server for the first time
# and need to accept key
if key_accept:
log.info('Adding {0} to known_hosts'.format(host))
self.conn.sendline('yes')
continue
else:
self.conn.sendline('no')
elif stdout and self.prompt_re.search(stdout):
# Auth success!
# We now have a prompt
break
def sendline(self, cmd):
'''
Send this command to the server and
return a tuple of the output and the stderr.
The format for parameters is:
cmd (string): The command to send to the sever.
'''
self.conn.sendline(cmd, self.linesep)
# saw_prompt = False
ret_stdout = []
ret_stderr = []
while self.conn.has_unread_data:
stdout, stderr = self.conn.recv()
if stdout:
ret_stdout.append(stdout)
if stderr:
log.debug('Error while executing command.')
ret_stderr.append(stderr)
if stdout and self.prompt_re.search(stdout):
break
return ''.join(ret_stdout), ''.join(ret_stderr)
def close_connection(self):
'''
Close the server connection
'''
self.conn.close(terminate=True, kill=True)