comware backup

routerbackup

@@ -325,7 +325,37 @@ class IosConfig(Config):
         return prepared
 
 
-class AsaConfig(Config):
+class ConfigWithShell(Config):
+    def _shell_command(self, channel, cmd,
+            waitfor=None,
+            waitmax=None,
+            send_timeout=10,
+            recv_timeout=20
+            ):
+        start = time.time()
+        while not channel.send_ready():
+            if time.time() - start > send_timeout:
+                raise RuntimeError('shell send timeout')
+            time.sleep(0.05)
+        channel.send(cmd +'\n')
+
+        start = time.time()
+        output = ''
+        while True:
+            while not channel.recv_ready():
+                if len(output) > 0:
+                    if waitfor and re.search(waitfor, output):
+                        return output.splitlines()
+                    if waitmax and time.time() - start > waitmax:
+                        return output.splitlines()
+                if time.time() - start > recv_timeout:
+                    raise RuntimeError(f'shell recv timeout (after receiving {len(output)} characters)')
+                time.sleep(0.05)
+            while channel.recv_ready():
+                output += str(channel.recv(65536), 'utf8')
+
+
+class AsaConfig(ConfigWithShell):
     def __init__(self, workdir):
         super().__init__(workdir)
         self.extra['iproute'] = {
@@ -333,33 +363,16 @@ class AsaConfig(Config):
             'content': None
             }
 
-    def asa_command(self, channel, cmd):
-        start = time.time()
-        while not channel.send_ready():
-            if time.time() - start > 10:
-                raise RuntimeError('ASA timeout')
-            time.sleep(0.05)
-        channel.send(cmd +'\n')
-        time.sleep(1)
-        while not channel.recv_ready():
-            if time.time() - start > 30:
-                raise RuntimeError('ASA timeout')
-            time.sleep(0.05)
-        output = ''
-        while channel.recv_ready():
-            output += str(channel.recv(65536), 'utf8')
-        return output.splitlines()
-
     def get_config(self):
         shrun = []
         channel = self.ssh.invoke_shell()
-        self.asa_command(channel, 'term pager 0')
-        shrun = self.asa_command(channel, 'sh run')
+        self._shell_command(channel, 'term pager 0', waitfor='#\s*$', waitmax=1)
+        shrun = self._shell_command(channel, 'sh run', waitfor=': end')
 
         iproute = []
         for cmd in ('sh ip address', 'sh route'):
             iproute.append('# '+ cmd.center(76, '='))
-            output = self.asa_command(channel, cmd)
+            output = self._shell_command(channel, cmd, waitfor='#\s*$', waitmax=2)
             iproute += output
 
         self.config = '\n'.join(shrun)
@@ -371,6 +384,22 @@ class AsaConfig(Config):
             return "does not contain ': end'."
         return False
 
+
+class ComwareConfig(ConfigWithShell):
+    def get_config(self):
+        shrun = []
+        channel = self.ssh.invoke_shell()
+        self._shell_command(channel, 'screen-length disable', waitfor='>\s*$', waitmax=1)
+        dicur = self._shell_command(channel, 'display current-configuration', waitfor='(?m:^return\s*$)')
+        self.config = '\n'.join(dicur)
+
+    def is_bad_config(self):
+        super().is_bad_config()
+        if not re.search('return', self.config):
+            return "does not contain 'return'."
+        return False
+
+
 def setup_logging(hostname, logfile, loglevel):
     if logfile:
         handler = logging.FileHandler(logfile, 'a')
@@ -395,7 +424,7 @@ def main():
     args = get_args()
     setup_logging(args.host, args.logfile, args.loglevel)
     try:
-        if args.type not in ('routeros', 'ios', 'asa'):
+        if args.type not in ('routeros', 'ios', 'asa', 'comware'):
             raise ValueError(f'Invalid devtype: {args.type}')
 
         if args.type == 'routeros':
@@ -404,6 +433,8 @@ def main():
             cf = IosConfig(workdir=args.backupdir)
         elif args.type == 'asa':
             cf = AsaConfig(workdir=args.backupdir)
+        elif args.type == 'comware':
+            cf = ComwareConfig(workdir=args.backupdir)
 
         address = args.address or args.host
         log.info(f'Connecting to {args.type}: {args.user}@{address} with{" password" if args.password else""}{" ssh-key" if args.sshkey else""}')