544 lines
21 KiB
Python
Executable File
544 lines
21 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
#Imports
|
|
import os
|
|
import re
|
|
import pexpect
|
|
from Crypto.PublicKey import RSA
|
|
from Crypto.Cipher import PKCS1_OAEP
|
|
import ast
|
|
from time import sleep
|
|
import datetime
|
|
import sys
|
|
import threading
|
|
|
|
#functions and classes
|
|
|
|
class node:
|
|
''' This class generates a node object. Containts all the information and methods to connect and interact with a device using ssh or telnet.
|
|
|
|
### Attributes:
|
|
|
|
- output (str): Output of the commands you ran with run or test
|
|
method.
|
|
|
|
- result(bool): True if expected value is found after running
|
|
the commands using test method.
|
|
'''
|
|
|
|
def __init__(self, unique, host, options='', logs='', password='', port='', protocol='', user='', config=''):
|
|
'''
|
|
|
|
### Parameters:
|
|
|
|
- unique (str): Unique name to assign to the node.
|
|
|
|
- host (str): IP address or hostname of the node.
|
|
|
|
### Optional Parameters:
|
|
|
|
- options (str): Additional options to pass the ssh/telnet for
|
|
connection.
|
|
|
|
- logs (str): Path/file for storing the logs. You can use
|
|
${unique},${host}, ${port}, ${user}, ${protocol}
|
|
as variables.
|
|
|
|
- password (str): Encrypted or plaintext password.
|
|
|
|
- port (str): Port to connect to node, default 22 for ssh and 23
|
|
for telnet.
|
|
|
|
- protocol (str): Select ssh or telnet. Default is ssh.
|
|
|
|
- user (str): Username to of the node.
|
|
|
|
- config (obj): Pass the object created with class configfile with
|
|
key for decryption and extra configuration if you
|
|
are using connection manager.
|
|
'''
|
|
if config == '':
|
|
self.idletime = 0
|
|
self.key = None
|
|
else:
|
|
self.idletime = config.config["idletime"]
|
|
self.key = config.key
|
|
self.unique = unique
|
|
attr = {"host": host, "logs": logs, "options":options, "port": port, "protocol": protocol, "user": user}
|
|
for key in attr:
|
|
profile = re.search("^@(.*)", attr[key])
|
|
if profile and config != '':
|
|
setattr(self,key,config.profiles[profile.group(1)][key])
|
|
elif attr[key] == '' and key == "protocol":
|
|
try:
|
|
setattr(self,key,config.profiles["default"][key])
|
|
except:
|
|
setattr(self,key,"ssh")
|
|
else:
|
|
setattr(self,key,attr[key])
|
|
if isinstance(password,list):
|
|
self.password = []
|
|
for i, s in enumerate(password):
|
|
profile = re.search("^@(.*)", password[i])
|
|
if profile and config != '':
|
|
self.password.append(config.profiles[profile.group(1)]["password"])
|
|
else:
|
|
self.password = [password]
|
|
|
|
def __passtx(self, passwords, *, keyfile=None):
|
|
# decrypts passwords, used by other methdos.
|
|
dpass = []
|
|
if keyfile is None:
|
|
keyfile = self.key
|
|
if keyfile is not None:
|
|
key = RSA.import_key(open(keyfile).read())
|
|
decryptor = PKCS1_OAEP.new(key)
|
|
for passwd in passwords:
|
|
if not re.match('^b[\"\'].+[\"\']$', passwd):
|
|
dpass.append(passwd)
|
|
else:
|
|
try:
|
|
decrypted = decryptor.decrypt(ast.literal_eval(passwd)).decode("utf-8")
|
|
dpass.append(decrypted)
|
|
except:
|
|
raise ValueError("Missing or corrupted key")
|
|
return dpass
|
|
|
|
|
|
|
|
def _logfile(self, logfile = None):
|
|
# translate logs variables and generate logs path.
|
|
if logfile == None:
|
|
logfile = self.logs
|
|
logfile = logfile.replace("${unique}", self.unique)
|
|
logfile = logfile.replace("${host}", self.host)
|
|
logfile = logfile.replace("${port}", self.port)
|
|
logfile = logfile.replace("${user}", self.user)
|
|
logfile = logfile.replace("${protocol}", self.protocol)
|
|
now = datetime.datetime.now()
|
|
dateconf = re.search(r'\$\{date \'(.*)\'}', logfile)
|
|
if dateconf:
|
|
logfile = re.sub(r'\$\{date (.*)}',now.strftime(dateconf.group(1)), logfile)
|
|
return logfile
|
|
|
|
def _logclean(self, logfile, var = False):
|
|
#Remove special ascii characters and other stuff from logfile.
|
|
if var == False:
|
|
t = open(logfile, "r").read()
|
|
else:
|
|
t = logfile
|
|
t = t.replace("\n","",1).replace("\a","")
|
|
t = t.replace('\n\n', '\n')
|
|
t = re.sub(r'.\[K', '', t)
|
|
while True:
|
|
tb = re.sub('.\b', '', t, count=1)
|
|
if len(t) == len(tb):
|
|
break
|
|
t = tb
|
|
ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/ ]*[@-~])')
|
|
t = ansi_escape.sub('', t)
|
|
if var == False:
|
|
d = open(logfile, "w")
|
|
d.write(t)
|
|
d.close()
|
|
return
|
|
else:
|
|
return t
|
|
|
|
def interact(self, debug = False):
|
|
'''
|
|
Allow user to interact with the node directly, mostly used by connection manager.
|
|
|
|
### Optional Parameters:
|
|
|
|
- debug (bool): If True, display all the connecting information
|
|
before interact. Default False.
|
|
'''
|
|
connect = self._connect(debug = debug)
|
|
if connect == True:
|
|
size = re.search('columns=([0-9]+).*lines=([0-9]+)',str(os.get_terminal_size()))
|
|
self.child.setwinsize(int(size.group(2)),int(size.group(1)))
|
|
print("Connected to " + self.unique + " at " + self.host + (":" if self.port != '' else '') + self.port + " via: " + self.protocol)
|
|
if 'logfile' in dir(self):
|
|
self.child.logfile_read = open(self.logfile, "wb")
|
|
elif debug:
|
|
self.child.logfile_read = None
|
|
if 'missingtext' in dir(self):
|
|
print(self.child.after.decode(), end='')
|
|
self.child.interact()
|
|
if "logfile" in dir(self) and not debug:
|
|
self._logclean(self.logfile)
|
|
else:
|
|
print(connect)
|
|
exit(1)
|
|
|
|
def run(self, commands,*, folder = '', prompt = r'>$|#$|\$$|>.$|#.$|\$.$', stdout = False):
|
|
'''
|
|
Run a command or list of commands on the node and return the output.
|
|
|
|
### Parameters:
|
|
|
|
- commands (str/list): Commands to run on the node. Should be
|
|
str or a list of str.
|
|
|
|
### Optional Named Parameters:
|
|
|
|
- folder (str): Path where output log should be stored, leave
|
|
empty to disable logging.
|
|
|
|
- prompt (str): Prompt to be expected after a command is finished
|
|
running. Usually linux uses ">" or EOF while
|
|
routers use ">" or "#". The default value should
|
|
work for most nodes. Change it if your connection
|
|
need some special symbol.
|
|
|
|
- stdout (bool):Set True to send the command output to stdout.
|
|
default False.
|
|
|
|
### Returns:
|
|
|
|
str: Output of the commands you ran on the node.
|
|
|
|
'''
|
|
connect = self._connect()
|
|
if connect == True:
|
|
expects = [prompt, pexpect.EOF]
|
|
output = ''
|
|
if isinstance(commands, list):
|
|
for c in commands:
|
|
result = self.child.expect(expects)
|
|
self.child.sendline(c)
|
|
if result == 0:
|
|
output = output + self.child.before.decode() + self.child.after.decode()
|
|
if result == 1:
|
|
output = output + self.child.before.decode()
|
|
else:
|
|
result = self.child.expect(expects)
|
|
self.child.sendline(commands)
|
|
if result == 0:
|
|
output = output + self.child.before.decode() + self.child.after.decode()
|
|
if result == 1:
|
|
output = output + self.child.before.decode()
|
|
result = self.child.expect(expects)
|
|
if result == 0:
|
|
output = output + self.child.before.decode() + self.child.after.decode()
|
|
if result == 1:
|
|
output = output + self.child.before.decode()
|
|
self.child.close()
|
|
output = output.lstrip()
|
|
if stdout == True:
|
|
print(output)
|
|
if folder != '':
|
|
with open(folder + "/" + self.unique, "w") as f:
|
|
f.write(output)
|
|
f.close()
|
|
self._logclean(folder + "/" + self.unique)
|
|
self.output = output
|
|
return output
|
|
else:
|
|
self.output = connect
|
|
return connect
|
|
|
|
def test(self, commands, expected, *, prompt = r'>$|#$|\$$|>.$|#.$|\$.$'):
|
|
'''
|
|
Run a command or list of commands on the node, then check if expected value appears on the output after the last command.
|
|
|
|
### Parameters:
|
|
|
|
- commands (str/list): Commands to run on the node. Should be
|
|
str or list of str.
|
|
|
|
- expected (str) : Expected text to appear after running
|
|
all the commands on the node.
|
|
|
|
### Optional Named Parameters:
|
|
|
|
- prompt (str): Prompt to be expected after a command is finished
|
|
running. Usually linux uses ">" or EOF while
|
|
routers use ">" or "#". The default value should
|
|
work for most nodes. Change it if your connection
|
|
need some special symbol.
|
|
|
|
### Returns:
|
|
bool: true if expected value is found after running the commands
|
|
false if prompt is found before.
|
|
|
|
'''
|
|
connect = self._connect()
|
|
if connect == True:
|
|
expects = [prompt, pexpect.EOF]
|
|
output = ''
|
|
if isinstance(commands, list):
|
|
for c in commands:
|
|
result = self.child.expect(expects)
|
|
self.child.sendline(c)
|
|
if result == 0:
|
|
output = output + self.child.before.decode() + self.child.after.decode()
|
|
if result == 1:
|
|
output = output + self.child.before.decode()
|
|
else:
|
|
self.child.expect(expects)
|
|
self.child.sendline(commands)
|
|
output = output + self.child.before.decode() + self.child.after.decode()
|
|
expects = [expected, prompt, pexpect.EOF]
|
|
results = self.child.expect(expects)
|
|
if results == 0:
|
|
self.child.close()
|
|
self.result = True
|
|
output = output + self.child.before.decode() + self.child.after.decode()
|
|
output = output.lstrip()
|
|
self.output = output
|
|
return True
|
|
if results in [1, 2]:
|
|
self.child.close()
|
|
self.result = False
|
|
if results == 1:
|
|
output = output + self.child.before.decode() + self.child.after.decode()
|
|
elif results == 2:
|
|
output = output + self.child.before.decode()
|
|
output = output.lstrip()
|
|
self.output = output
|
|
return False
|
|
else:
|
|
self.result = None
|
|
self.output = connect
|
|
return connect
|
|
|
|
def _connect(self, debug = False):
|
|
# Method to connect to the node, it parse all the information, create the ssh/telnet command and login to the node.
|
|
if self.protocol == "ssh":
|
|
cmd = "ssh"
|
|
if self.idletime > 0:
|
|
cmd = cmd + " -o ServerAliveInterval=" + str(self.idletime)
|
|
if self.user == '':
|
|
cmd = cmd + " -t {}".format(self.host)
|
|
else:
|
|
cmd = cmd + " -t {}".format("@".join([self.user,self.host]))
|
|
if self.port != '':
|
|
cmd = cmd + " -p " + self.port
|
|
if self.options != '':
|
|
cmd = cmd + " " + self.options
|
|
if self.logs != '':
|
|
self.logfile = self._logfile()
|
|
if self.password[0] != '':
|
|
passwords = self.__passtx(self.password)
|
|
else:
|
|
passwords = []
|
|
expects = ['yes/no', 'refused', 'supported', 'cipher', 'sage', 'timeout', 'unavailable', 'closed', '[p|P]assword:|[u|U]sername:', r'>$|#$|\$$|>.$|#.$|\$.$', 'suspend', pexpect.EOF, "No route to host", "resolve hostname", "no matching host key"]
|
|
elif self.protocol == "telnet":
|
|
cmd = "telnet " + self.host
|
|
if self.port != '':
|
|
cmd = cmd + " " + self.port
|
|
if self.options != '':
|
|
cmd = cmd + " " + self.options
|
|
if self.logs != '':
|
|
self.logfile = self._logfile()
|
|
if self.password[0] != '':
|
|
passwords = self.__passtx(self.password)
|
|
else:
|
|
passwords = []
|
|
expects = ['[u|U]sername:', 'refused', 'supported', 'cipher', 'sage', 'timeout', 'unavailable', 'closed', '[p|P]assword:', r'>$|#$|\$$|>.$|#.$|\$.$', 'suspend', pexpect.EOF, "No route to host", "resolve hostname", "no matching host key"]
|
|
else:
|
|
raise ValueError("Invalid protocol: " + self.protocol)
|
|
child = pexpect.spawn(cmd)
|
|
if debug:
|
|
child.logfile_read = sys.stdout.buffer
|
|
if len(passwords) > 0:
|
|
loops = len(passwords)
|
|
else:
|
|
loops = 1
|
|
endloop = False
|
|
for i in range(0, loops):
|
|
while True:
|
|
results = child.expect(expects)
|
|
if results == 0:
|
|
if self.protocol == "ssh":
|
|
child.sendline('yes')
|
|
elif self.protocol == "telnet":
|
|
if self.user != '':
|
|
child.sendline(self.user)
|
|
else:
|
|
self.missingtext = True
|
|
break
|
|
if results in [1, 2, 3, 4, 5, 6, 7, 12, 13, 14]:
|
|
child.close()
|
|
return "Connection failed code:" + str(results)
|
|
if results == 8:
|
|
if len(passwords) > 0:
|
|
child.sendline(passwords[i])
|
|
else:
|
|
self.missingtext = True
|
|
break
|
|
if results in [9, 11]:
|
|
endloop = True
|
|
child.sendline()
|
|
break
|
|
if results == 10:
|
|
child.sendline("\r")
|
|
sleep(2)
|
|
if endloop:
|
|
break
|
|
child.readline(0)
|
|
self.child = child
|
|
return True
|
|
|
|
class nodes:
|
|
''' This class generates a nodes object. Contains a list of node class objects and methods to run multiple tasks on nodes simultaneously.
|
|
|
|
### Attributes:
|
|
|
|
- nodelist (list): List of node class objects passed to the init
|
|
function.
|
|
|
|
- output (dict): Dictionary formed by nodes unique as keys,
|
|
output of the commands you ran on the node as
|
|
value. Created after running methods run or test.
|
|
|
|
- result (dict): Dictionary formed by nodes unique as keys, value
|
|
is True if expected value is found after running
|
|
the commands, False if prompt is found before.
|
|
Created after running method test.
|
|
|
|
- <unique> (obj): For each item in nodelist, there is an attribute
|
|
generated with the node unique.
|
|
'''
|
|
|
|
def __init__(self, nodes: dict, config = ''):
|
|
'''
|
|
### Parameters:
|
|
|
|
- nodes (dict): Dictionary formed by node information:
|
|
Keys: Unique name for each node.
|
|
Mandatory Subkeys: host(str).
|
|
Optional Subkeys: options(str), logs(str), password(str),
|
|
port(str), protocol(str), user(str).
|
|
For reference on subkeys check node class.
|
|
|
|
### Optional Parameters:
|
|
|
|
- config (obj): Pass the object created with class configfile with key
|
|
for decryption and extra configuration if you are using
|
|
connection manager.
|
|
'''
|
|
self.nodelist = []
|
|
self.config = config
|
|
for n in nodes:
|
|
this = node(n, **nodes[n], config = config)
|
|
self.nodelist.append(this)
|
|
setattr(self,n,this)
|
|
|
|
|
|
def _splitlist(self, lst, n):
|
|
#split a list in lists of n members.
|
|
for i in range(0, len(lst), n):
|
|
yield lst[i:i + n]
|
|
|
|
|
|
def run(self, commands,*, folder = None, prompt = None, stdout = None, parallel = 10):
|
|
'''
|
|
Run a command or list of commands on all the nodes in nodelist.
|
|
|
|
### Parameters:
|
|
|
|
commands (str/list): Commands to run on the node. Should be str or
|
|
list of str.
|
|
|
|
### Optional Named Parameters:
|
|
|
|
folder (str): Path where output log should be stored, leave empty
|
|
to disable logging.
|
|
|
|
prompt (str): Prompt to be expected after a command is finished
|
|
running. Usually linux uses ">" or EOF while routers
|
|
use ">" or "#". The default value should work for
|
|
most nodes. Change it if your connection need some
|
|
special symbol.
|
|
|
|
stdout (bool): Set True to send the command output to stdout.
|
|
Default False.
|
|
|
|
parallel (int): Number of nodes to run the commands simultaneously.
|
|
Default is 10, if there are more nodes that this
|
|
value, nodes are groups in groups with max this
|
|
number of members.
|
|
|
|
###Returns:
|
|
|
|
dict: Dictionary formed by nodes unique as keys, Output of the
|
|
commands you ran on the node as value.
|
|
|
|
'''
|
|
args = {}
|
|
args["commands"] = commands
|
|
if folder != None:
|
|
args["folder"] = folder
|
|
if prompt != None:
|
|
args["prompt"] = prompt
|
|
if stdout != None:
|
|
args["stdout"] = stdout
|
|
output = {}
|
|
tasks = []
|
|
for n in self.nodelist:
|
|
tasks.append(threading.Thread(target=n.run, kwargs=args))
|
|
taskslist = list(self._splitlist(tasks, parallel))
|
|
for t in taskslist:
|
|
for i in t:
|
|
i.start()
|
|
for i in t:
|
|
i.join()
|
|
for i in self.nodelist:
|
|
output[i.unique] = i.output
|
|
self.output = output
|
|
return output
|
|
|
|
def test(self, commands, expected, *, prompt = None, parallel = 10):
|
|
'''
|
|
Run a command or list of commands on all the nodes in nodelist, then check if expected value appears on the output after the last command.
|
|
|
|
### Parameters:
|
|
|
|
commands (str/list): Commands to run on the node. Should be str or
|
|
list of str.
|
|
|
|
expected (str) : Expected text to appear after running all the
|
|
commands on the node.
|
|
|
|
### Optional Named Parameters:
|
|
|
|
prompt (str): Prompt to be expected after a command is finished
|
|
running. Usually linux uses ">" or EOF while
|
|
routers use ">" or "#". The default value should
|
|
work for most nodes. Change it if your connection
|
|
need some special symbol.
|
|
|
|
### Returns:
|
|
|
|
dict: Dictionary formed by nodes unique as keys, value is True if
|
|
expected value is found after running the commands, False
|
|
if prompt is found before.
|
|
|
|
'''
|
|
args = {}
|
|
args["commands"] = commands
|
|
args["expected"] = expected
|
|
if prompt != None:
|
|
args["prompt"] = prompt
|
|
output = {}
|
|
result = {}
|
|
tasks = []
|
|
for n in self.nodelist:
|
|
tasks.append(threading.Thread(target=n.test, kwargs=args))
|
|
taskslist = list(self._splitlist(tasks, parallel))
|
|
for t in taskslist:
|
|
for i in t:
|
|
i.start()
|
|
for i in t:
|
|
i.join()
|
|
for i in self.nodelist:
|
|
result[i.unique] = i.result
|
|
output[i.unique] = i.output
|
|
self.output = output
|
|
self.result = result
|
|
return result
|
|
|
|
# script
|