D7net
Home
Console
Upload
information
Create File
Create Folder
About
Tools
:
/
proc
/
2
/
task
/
2
/
root
/
opt
/
imh-python
/
lib
/
python2.7
/
site-packages
/
cmu_ded
/
Libs
/
Filename :
tbSSH.py
back
Copy
""" The idea is to setup a default class for ssh specifically for IMH. """ import os import re import sys import subprocess import socket import getpass from logthis import LogThis from Crypto.PublicKey import RSA from subprocess import CalledProcessError from paramiko import SSHClient, AutoAddPolicy, ssh_exception, RSAKey, Transport, SSHException, SFTPClient class TbSSH(LogThis): def __init__(self, sup_user, host, supplied_port, full_connect=False, sup_password=None, sup_sshkey=None, logger=None, rsyncspeed=10000): """ Default constructor for TbSSH Not everything works in here. I am porting this over from a different project. :param sup_user: Simply the user you want to conenct with :param host: The Host you want to connect with :param supplied_port: The port you want to connect over. :param full_connect: This is the part where is works with IMH. Full connects check's the host name of the server then will execute the function needed to connect. :param sup_password: The password needed to get in. :param sup_sshkey: The SSH_key needed to get in. This can't be raw text and has to be a pkey object. To create :param logger: logthis.py object. Used for printing Logging info. this object it needs a file, or a file like object. From there we use RSAKey to create the pkey object then pass it through to paramiko to connect. http://docs.paramiko.org/en/2.1/api/keys.html#paramiko.pkey.PKey """ super(TbSSH, self).__init__(logger=logger) self.t = None self.sftp = None self.client = SSHClient() self.channel = None self.sup_pkey = None self.gen_sshkey = False self.gen_pkey = "" self.user = sup_user self.host = host self.port = supplied_port self.password = sup_password self.sshkey = sup_sshkey self.rsyncspeed = rsyncspeed # Init logging. self.from_server = self.host self.to_server = socket.gethostname() if full_connect: # Not setup yet # Setting up patterns. vps_pattern = re.compile('vps') ded_pattern = re.compile('ded|dedicated|essential|advanced|elite|cc') # Connecting to shared. if vps_pattern.match(host): self.vps_connect(host) elif ded_pattern.match(host): self.ded_connect(host) else: self.full_ssh_connection(host, supplied_port, sup_user) else: # Trying to auth 3 times. for _ in range(3): if self.attempt_connection(): self.channel = self.client.invoke_shell() break else: self.info("Trying again....") if not self.t.is_authenticated(): self.critical("It appears we could not make a connection. " "This is generally due to an issue with the SSHD configurations.\n" "Validate the password or ssh-key is correct, " "otherwise validate things like permit root ssh login are set to true.") # Close every time def close(self): """ Closing the connection. Also removing temp keys. :return: Closed SSH connection. """ if self.gen_pkey: self.remove_authoization() os.remove(self.gen_pkey) if self.gen_sshkey: os.remove(self.sshkey) self.client.close() def attempt_connection(self): """ Attempting connection to the server. :return: Returns the connection. """ self.client.set_missing_host_key_policy(AutoAddPolicy()) self.client.load_system_host_keys() # Trying to generate a transport first. try: self.t = Transport((self.host, self.port)) self.t.start_client() except SSHException as e: self.error("Connection issues. Please validate the IP is added to the firewall. Below is the error.") self.error(e.message) sys.exit(1) # Segmentation. if self.sshkey is not None: result = self.sshkey_auth() if not result: return result else: result = self.user_pass_auth(self.t) if not result: return result self.sftp = SFTPClient.from_transport(self.t) return True def user_pass_auth(self, transport): """ Auth setting for user/pass combo. The __init__ function loops a couple times to get the correct connection. :return: Returns true if connection is successful. Resets self.password if it fails. """ try: self.t.auth_password(self.user, self.password) except ssh_exception.BadAuthenticationType as e: self.error("Bad Authentication type. This is generally a sign that permit root login is False") self.error("Below are the following allowed types for authentication. " "Validate password is in the list") self.error(str(e.allowed_types)) except ssh_exception.AuthenticationException: self.password = getpass.getpass("[x] Authentication Failed....\nPlease input the password: ") return False try: self.client.connect(self.host, username=self.user, password=self.password, port=self.port) except ssh_exception.AuthenticationException: self.password = getpass.getpass("[x] Authentication Failed....\nPlease input the password: ") return False self.gen_rsa_key() self.authorize_key() return True def sshkey_auth(self): """ SSHKey connection. The __init__ function loops a couple times till the connection is established. :return: True if connection. False if something has gone wrong. The problem item is reset before. """ # Cleaning up error reporting by setting sup_pkey to None first. try: self.sup_pkey = RSAKey.from_private_key(file_obj=open(self.sshkey), password=self.password) except ssh_exception.PasswordRequiredException: self.password = getpass.getpass(prompt="[X] RSAKey has a password on it...\nInput its password for me: ") return False try: self.t.auth_publickey(self.user, self.sup_pkey) except ssh_exception.BadAuthenticationType as e: self.error("Bad Authentication type.") self.error("Below are the following allowed types for authentication. Validate publickey is in the list") self.error(str(e.allowed_types)) sys.exit(-1) except ssh_exception.AuthenticationException as e: self.error(str(e.message)) self.error("Please add the ssh-key to the users authorized_keys.") sys.exit(-1) try: # pkey is being used instead of key_filename due to flexability with cPanel API. self.client.connect(self.host, port=self.port, username=self.user, pkey=self.sup_pkey) except ssh_exception.AuthenticationException: self.password = getpass.getpass(prompt="[x] Authentication Failed....\nPlease input the password: ") return False pass return True def ssh_exec(self, command): """ Executing the command. Setting get_pty in the hopes of helping resolve utmp errors. :param command: Command you want to run. :return: Returns a tuple of (stdin, stdout, stderr) """ stdin, stdout, stderr = self.client.exec_command(command, get_pty=True) return stdin, stdout, stderr def string_exec(self, command): """ Simple idea, simply return the string of the command ran. Uses stdout and reads from that. :param command: The command you want to run :return: Return the string. """ output = self.ssh_exec(command)[1].readlines() return "".join(output) def print_exec(self, command): """ Simply printing the command. string_exec makes this easy. :param command: Command you want to run. :return: Nothing. (NULLLLLLL) """ print self.string_exec(command) # Screen -R either creates the screen or attaches. def attach_screen(self, screen_name): """ Attaching a screen. Using -R as it either attachs a screen or creates one. :param screen_name: Screen name you want to attach :return: Nothing. """ self.ssh_exec('screen -R %s' % screen_name) # Sending ^A + d def detach_screen(self): """ Detaching screen. This is done by sending ctrl + A then d :return: nothing """ self.channel.send('\x01d') def vps_connect(self, server): print('VPS conect Not setup yet') def ded_connect(self, server): print('Ded connect Not setup yet') def full_ssh_connection(self, server, port, user='root'): """ Full on SSH conection with the server. This is used if you want to do things on the server with the connection. This useses paramiko's communicate function to take over your shell. :param server: Server you want to connect to. :param port: The port of the server :param user: The username. :return: Nothing, this just uses """ terminal = subprocess.Popen(['/usr/bin/ssh', '-o', 'StrictHostKeyChecking=no', '-o', 'UserKnownHostsFile=/dev/null', '-p', str(port), '%s@%s' % (user, server)]) terminal.communicate() def file_exists(self, path): """ Trying to see if the path exists through sftp. :param path: Path of the file you want to check. :return: True if path exists. """ try: self.sftp.stat(path) except IOError, e: if e[0] == 2: return False raise else: return True def user_exists(self, user): """ Running id -u on the remote server to validate the user is there. This returns the ID of the user or no such user if there is no user. :param user: user you want to check. :return: True if the user exists. False otherwise. """ if 'no such user' in self.string_exec('id -u %s' % user): return False else: return True def pull_files(self, remote_dir, file_to_transfer=None, local_dir='./', recursive=False): """ Function that grabs files. Basically, it checks each file. If file_to_transfer is none, then everything is transfered. If not none, It checks the name of the file then transfers. For example, if cpmove is passed, it will only transfer cpmove files. From there, it checks if a file is a symlink or a directory. If its a directory, and recursive is true, we call pull_files again but in the directory. If its a symlink, ignore it. ( For now ) :param remote_dir: The remote directory :param file_to_transfer: Specific names to transfer :param local_dir: The local directory :param recursive: If you want to recursivly download everything. :return: And array of errors. File it failed on then the exception raised. """ result = [] remote_dir = remote_dir + '/' local_dir = local_dir + '/' self.sftp.chdir(remote_dir) for file_name in self.sftp.listdir(): try: if file_to_transfer is not None and file_to_transfer not in file_name: continue filestats = str(self.sftp.lstat(file_name)).split()[0] # checking attributes of file to see if it is a directory. # Or symlink. if 'l' in filestats: continue if 'd' in filestats: if recursive: os.mkdir(local_dir + file_name) self.pull_files(remote_dir + file_name, file_to_transfer=file_to_transfer, local_dir=local_dir + file_name, recursive=recursive) self.sftp.chdir(remote_dir) else: continue else: self.debug("sftpget is being called for %s" % file_name) self.sftp.get(file_name, local_dir + file_name) except Exception as e: result.append(file_name) result.append(e) pass return result def rsync_pull_files(self, directory_to_grab, directory_to_store): """ Rsyncing files from one server to another. This was created as the pull files function above was really slow. rsync is here if we need to move files faster. sshkeys need to be setup first. :param directory_to_grab: Directory your pulling. :param directory_to_store: Directory your storing. :return: Returns all errors. """ fnull = open(os.devnull, 'w') command = [ '/bin/rsync', '-a', '--quiet', '--rsh=ssh -i %s -o StrictHostKeyChecking=no -p %s' % (self.sshkey, self.port), '--bwlimit=%d' % self.rsyncspeed ] command.append('%s@%s:%s' % (self.user, self.host, directory_to_grab)) command.append('%s' % directory_to_store) # if rsyncspeed is under 0, then we remove it from the command. if self.rsyncspeed <= 0: command.pop(4) # removing errors for now. If you need to debug, turn them back on. try: subprocess.check_call(command, stderr=fnull) if os.path.exists("/opt/cmu_ded/password"): os.remove("/opt/cmu_ded/password") except CalledProcessError as e: self.error("Whops, Rsync broke.") self.error("Rsync recieved the following error code: " + str(e.returncode)) if os.path.exists("/opt/cmu_ded/password"): os.remove("/opt/cmu_ded/password") return e.returncode pass def gen_rsa_key(self): """ Setter for self.key_file and self.public_key_file :return: Nothing. """ self.info("Generating ssh-key") self.gen_sshkey = True key = RSA.generate(2048) with open("/opt/cmu_ded/tmp/private.key", 'w') as content_file: os.chmod("/opt/cmu_ded/tmp/private.key", 0600) self.sshkey = "/opt/cmu_ded/tmp/private.key" content_file.write(key.exportKey('PEM')) pubkey = key.publickey() with open("/opt/cmu_ded/tmp/public.key", 'w') as content_file: content_file.write(pubkey.exportKey('OpenSSH')) self.gen_pkey = "/opt/cmu_ded/tmp/public.key" def authorize_key(self): """ Adds key to authorized hosts on the remote server. Need to check if .ssh exists. :return: Nothing. """ self.info("Authorizing Key") public_key_contents = open(self.gen_pkey).read() # If the directory doesn't exist, create it. result = self.string_exec("stat /root/.ssh") if "No such file or directory" in result: self.string_exec("mkdir /root/.ssh") self.string_exec("chown root:root /root/.ssh") self.string_exec("chmod 700 /root/.ssh") return self.string_exec("echo %s >> /root/.ssh/authorized_keys" % public_key_contents) def remove_authoization(self): """ So we don't blow up a authorized_key file. :return: Removes SSH key. """ self.info("Removing Authorized Key") public_key_contents = open(self.gen_pkey).read() # Bash needs the value to be escaped at some level. Make sure to use re.escape() self.ssh_exec("/bin/grep -v %s /root/.ssh/authorized_keys > /root/.ssh/authorized_keys.old" % re.escape(public_key_contents)) self.ssh_exec("/bin/rm -f /root/.ssh/authorized_keys") self.ssh_exec("/bin/mv /root/.ssh/authorized_keys.old /root/.ssh/authorized_keys") def check_connection(self): try: self.ssh_exec("ls") return True except socket.error as e: # attempt to connect again. self.warning("[!!] SSH Connection has been closed. Attempting to regain the connection.") if self.attempt_connection(): self.info("[*] SSH Connection re-established.") else: self.error("!!! SSH Connection failed to reconnect.")