From 255b2bd4ef21c39fa64bf828d26755a562df3c12 Mon Sep 17 00:00:00 2001
From: fluzzi <fluzzi@gmail.com>
Date: Fri, 18 Mar 2022 15:32:48 -0300
Subject: [PATCH] update big

---
 conn/configfile.py |  79 +++++++++++++++++++++++++++++---
 conn/core.py       | 110 +++++++++++++++++++++++++++++++--------------
 test.py            |  28 ++++++++++--
 3 files changed, 173 insertions(+), 44 deletions(-)

diff --git a/conn/configfile.py b/conn/configfile.py
index 702a136..37834ca 100755
--- a/conn/configfile.py
+++ b/conn/configfile.py
@@ -15,7 +15,7 @@ from Crypto.PublicKey import RSA
 
 class configfile:
     
-    def __init__(self, conf = None):
+    def __init__(self, conf = None, *, key = None):
         home = os.path.expanduser("~")
         self.defaultdir = home + '/.config/conn'
         self.defaultfile = self.defaultdir + '/config.yaml'
@@ -24,8 +24,11 @@ class configfile:
             self.dir = self.defaultdir
             self.file = self.defaultfile
         else:
-            self.dir = os.path.dirname(conf)
             self.file = conf
+        if key == None:
+            self.key = self.defaultkey
+        else:
+            self.key = key
         if os.path.exists(self.file):
             config = self.loadconfig(self.file)
         else:
@@ -33,23 +36,85 @@ class configfile:
         self.config = config["config"]
         self.connections = config["connections"]
         self.profiles = config["profiles"]
+        if not os.path.exists(self.key):
+            self.createkey(self.key)
+        self.privatekey = RSA.import_key(open(self.key).read())
+        self.publickey = self.privatekey.publickey()
 
 
     def loadconfig(self, conf):
         ymlconf = open(conf)
-        return yaml.load(ymlconf.read(), Loader=yaml.FullLoader)
+        return yaml.load(ymlconf.read(), Loader=yaml.CLoader)
 
     def createconfig(self, conf):
-        defaultconfig = {'config': {'case': False, 'frun': False, 'idletime': 30}, 'connections': {}, 'profiles': {}}
+        defaultconfig = {'config': {'case': False, 'frun': False, 'idletime': 30}, 'connections': {}, 'profiles': { "default": { "host":"", "protocol":"ssh", "port":"", "user":"", "password":"", "options":"", "logs":"" }}}
         if not os.path.exists(conf):
             with open(conf, "w") as f:
-                yaml.dump(defaultconfig, f, explicit_start=True)
+                yaml.dump(defaultconfig, f, explicit_start=True, Dumper=yaml.CDumper)
                 f.close()
         ymlconf = open(conf)
-        return yaml.load(ymlconf.read(), Loader=yaml.FullLoader)
+        return yaml.load(ymlconf.read(), Loader=yaml.CLoader)
+
+    def saveconfig(self, conf):
+        newconfig = {"config":{}, "connections": {}, "profiles": {}}
+        newconfig["config"] = self.config
+        newconfig["connections"] = self.connections
+        newconfig["profiles"] = self.profiles
+        with open(conf, "w") as f:
+            yaml.dump(newconfig, f, explicit_start=True, Dumper=yaml.CDumper)
+            f.close()
 
     def createkey(self, keyfile):
         key = RSA.generate(2048)
-        with open('mykey.pem','wb') as f:
+        with open(keyfile,'wb') as f:
             f.write(key.export_key('PEM'))
             f.close()
+
+    def _explode_unique(self, unique):
+        uniques = unique.split("@")
+        result = {"id": uniques[0]}
+        if len(uniques) == 2:
+            result["folder"] = uniques[1]
+        elif len(uniques) == 3:
+            result["folder"] = uniques[2]
+            result["subfolder"] = uniques[1]
+        return result
+
+    def _connections_add(self,*, id, host, folder='', subfolder='', options='', logs='', password='', port='', protocol='', user='' ):
+        if folder == '':
+            self.connections[id] = {"host": host, "options": options, "logs": logs, "password": password, "port": port, "protocol": protocol, "user": user, "type": "connection"}
+        elif folder != '' and subfolder == '':
+            self.connections[folder][id] = {"host": host, "options": options, "logs": logs, "password": password, "port": port, "protocol": protocol, "user": user, "type": "connection"}
+        elif folder != '' and subfolder != '':
+            self.connections[folder][subfolder][id] = {"host": host, "options": options, "logs": logs, "password": password, "port": port, "protocol": protocol, "user": user, "type": "connection"}
+            
+
+    def _connections_del(self,*, id, folder='', subfolder=''):
+        if folder == '':
+            del self.connections[id]
+        elif folder != '' and subfolder == '':
+            del self.connections[folder][id]
+        elif folder != '' and subfolder != '':
+            del self.connections[folder][subfolder][id]
+
+    def _folder_add(self,*, folder, subfolder = ''):
+        if subfolder == '':
+            if folder not in self.connections:
+                self.connections[folder] = {"type": "folder"}
+        else:
+            if subfolder not in self.connections[folder]:
+                self.connections[folder][subfolder] = {"type": "subfolder"}
+
+    def _folder_del(self,*, folder, subfolder=''):
+        if subfolder == '':
+            del self.connections[folder]
+        else:
+            del self.connections[folder][subfolder]
+
+
+    def _profiles_add(self,*, id, host = '', options='', logs='', password='', port='', protocol='', user='' ):
+        self.profiles[id] = {"host": host, "options": options, "logs": logs, "password": password, "port": port, "protocol": protocol, "user": user}
+            
+
+    def _profiles_del(self,*, id ):
+        del self.profiles[id]
diff --git a/conn/core.py b/conn/core.py
index 97f04c6..13727da 100755
--- a/conn/core.py
+++ b/conn/core.py
@@ -20,19 +20,19 @@ from .configfile import configfile
 #functions and clsses
 
 class node:
-    def __init__(self, unique, host, options='', logs='', password='', port='', protocol='', type='', user=''):
-        try:
-            config = configfile()
-            self.idletime = config.config["idletime"]
-        except:
-            config = {}
+    def __init__(self, unique, host, options='', logs='', password='', port='', protocol='', type='', user='', config=''):
+        if config == '':
             self.idletime = 0
+            self.key = None
+        else:
+            self.idletime = config.config["idletime"]
+            self.key = config.key
         self.unique = unique
         self.id = self.unique.split("@")[0]
         attr = {"host": host, "logs": logs, "options":options, "port": port, "protocol": protocol, "user": user}
         for key in attr:
             profile = re.search("^@(.*)", attr[key])
-            if profile:
+            if profile and config != '':
                 setattr(self,key,config.profiles[profile.group(1)][key])
             elif attr[key] == '' and key == "protocol":
                 try:
@@ -45,23 +45,29 @@ class node:
             self.password = []
             for i, s in enumerate(password):
                 profile = re.search("^@(.*)", password[i])
-                if profile:
+                if profile and config != '':
                     self.password.append(config.profiles[profile.group(1)]["password"])
         else:
-            self.password = [str(password)]
+            self.password = [password]
 
     def __passtx(self, passwords, *, keyfile=None):
-        if keyfile is None:
-            home = os.path.expanduser("~")
-            keyfile = home + '/.config/conn/.osk'
-        key = RSA.import_key(open(keyfile).read())
-        publickey = key.publickey()
-        encryptor = PKCS1_OAEP.new(publickey)
-        decryptor = PKCS1_OAEP.new(key)
+        keyfile = self.key
         dpass = []
+        if keyfile is None:
+            keyfile = self.key
+        else:
+            key = RSA.import_key(open(keyfile).read())
+            decryptor = PKCS1_OAEP.new(key)
         for passwd in passwords:
-            decrypted = decryptor.decrypt(ast.literal_eval(str(passwd))).decode("utf-8")
-            dpass.append(decrypted)
+            if isinstance(passwd, str):
+                dpass.append(passwd)
+            else:
+                try:
+                    decrypted = decryptor.decrypt(ast.literal_eval(str(passwd))).decode("utf-8")
+                    dpass.append(decrypted)
+                except:
+                    print("Missing or wrong key")
+                    exit(1)
         return dpass
 
     
@@ -90,12 +96,53 @@ class node:
             if len(t) == len(tb):
                 break
             t = tb
+        ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/ ]*[@-~])')
+        t = ansi_escape.sub('', t)
         d = open(logfile, "w")
         d.write(t)
         d.close()
         return
 
-    def connect(self, mode = None):
+    def interact(self, missingtext = False):
+        connect = self._connect()
+        if connect == True:
+            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")
+                # self.child.logfile = sys.stdout.buffer
+            if 'missingtext' in dir(self):
+                print(child.after.decode(), end='')
+            self.child.interact()
+            if "logfile" in dir(self):
+                self._logclean(self.logfile)
+
+    def run(self, commands,*, folder = '', prompt = '>$|#$|\$.$'):
+        connect = self._connect()
+        if connect == True:
+            output = ''
+            if isinstance(commands, list):
+                for c in commands:
+                    self.child.expect(prompt)
+                    self.child.sendline(c)
+                    output = output + self.child.before.decode() + self.child.after.decode()
+            else:
+                self.child.expect(prompt)
+                print(self.child.sendline(commands))
+                output = output + self.child.before.decode() + self.child.after.decode()
+            self.child.expect(prompt)
+            output = output + self.child.before.decode() + self.child.after.decode()
+            if folder == '':
+                print(output)
+            else:
+                with open(folder + "/" + self.unique, "w") as f:
+                    f.write(output)
+                    f.close()
+                    self._logclean(folder + "/" + self.unique)
+
+
+            
+
+    def _connect(self):
         if self.protocol == "ssh":
             cmd = "ssh"
             if self.idletime > 0:
@@ -109,12 +156,12 @@ class node:
             if self.options != '':
                 cmd = cmd + " " + self.options
             if self.logs != '':
-                logfile = self._logfile()
+                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:', '>$|#$|\$', 'suspend']
+            expects = ['yes/no', 'refused', 'supported', 'cipher', 'sage', 'timeout', 'unavailable', 'closed', '[p|P]assword:|[u|U]sername:', '>$|#$|\$.$', 'suspend', pexpect.EOF]
         elif self.protocol == "telnet":
             cmd = "telnet " + self.host
             if self.port != '':
@@ -122,20 +169,16 @@ class node:
             if self.options != '':
                 cmd = cmd + " " + self.options
             if self.logs != '':
-                logfile = self._logfile()
+                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:', '>$|#$|\$', 'suspend']
+            expects = ['[u|U]sername:', 'refused', 'supported', 'cipher', 'sage', 'timeout', 'unavailable', 'closed', '[p|P]assword:', '>$|#$|\$.$', 'suspend', pexpect.EOF]
         else:
             print("Invalid protocol: " + self.protocol)
             return
         child = pexpect.spawn(cmd)
-        if 'logfile' in locals():
-            child.logfile_read = open(logfile, "wb")
-            # child.logfile_read = sys.stdout.buffer
-        print("Connecting to " + self.unique + " at " + self.host + (":" if self.port != '' else '') + self.port + " via: " + self.protocol)
         if len(passwords) > 0:
             loops = len(passwords)
         else:
@@ -152,7 +195,7 @@ class node:
                             if self.user != '':
                                 child.sendline(self.user)
                             else:
-                                print(child.after.decode(), end='')
+                                self.missingtext = True
                                 break
                     case 1 | 2 | 3 | 4 | 5 | 6 |7:
                         print("Connection failed code:" + str(results))
@@ -162,9 +205,9 @@ class node:
                         if len(passwords) > 0:
                             child.sendline(passwords[i])
                         else:
-                            print(child.after.decode(), end='')
+                            self.missingtext = True
                         break
-                    case 9:
+                    case 9 | 11:
                         endloop = True
                         child.sendline()
                         break
@@ -174,9 +217,8 @@ class node:
             if endloop:
                 break
         child.readline(0)
-        if mode == "interact":
-            child.interact()
-            if "logfile" in locals():
-                self._logclean(logfile)
+        self.child = child
+        return True
+
 
 # script
diff --git a/test.py b/test.py
index b0961dc..d9578bf 100755
--- a/test.py
+++ b/test.py
@@ -1,7 +1,29 @@
 #!/usr/bin/python3
 import conn
+import yaml
 
-conf = conn.configfile()
-conn1=conn.node("pruebas@conn", **conf.connections["home"]["xr"])
-conn1.connect("interact")
+conf = conn.configfile("test.yaml")
+# ***
+# conf._connections_del(id = "zab3mu", folder="teco")
+# conf._connections_add(id = "zzztest", folder="teco" ,host = "10.21.96.45", user="sarabada")
+# conf._folder_add(folder="zzz")
+# conf._folder_add(folder="zzz", subfolder="achus")
+# conf._connections_add(id = "zzztec", subfolder="achus", folder="zzz" ,host = "10.21.96.45")
+# conf._connections_add(id = "zzztec", subfolder="achus", folder="zzz" ,host = "10.21.96.45", options=" saracatanga")
+# conf._folder_del(folder = "zzz", subfolder = "achus")
+# conf._profiles_add(id = "test", user = 'tesuser')
+# conf._profiles_add(id = "test", user = 'tesuser', protocol = 'telnet')
+# conf._profiles_del(id = "test")
+# print(yaml.dump(conf.profiles))
+# conf.saveconfig("test.yaml")
+# ***
+xr=conn.node("xr@home", **conf.connections["home"]["xr"], config=conf)
+ios=conn.node("ios@home", **conf.connections["home"]["ios"], config=conf)
+norman = conn.node("norman@home", **conf.connections["home"]["norman"], config=conf)
+eve = conn.node("eve@home", **conf.connections["home"]["eve"], config=conf)
+# xr.run(["term len 0","show ip bgp", "show ip bgp summ"], folder="test")
+# ios.run(["term len 0","show ip bgp", "show ip bgp summ"], folder="test")
+# norman.run(["ls -la", "pwd"], folder = "test")
+# eve.run(["ls -la", "pwd"], folder = "test")
+xr.interact()