Merge pull request #4 from SamBriskman/Sam

Update test case

Author: vincentbergman

components/netdriver/plugins/check_point/check_point.py

@@ -37,7 +37,7 @@ def get_mode_prompt_patterns(self) -> dict[Mode, re.Pattern]:
         }
 
     def get_more_pattern(self) -> tuple[re.Pattern, str]:
-        return (CheckPointBase.PatternHelper.get_more_pattern, self._CMD_MORE)
+        return (CheckPointBase.PatternHelper.get_more_pattern(), self._CMD_MORE)
 
     class PatternHelper:
         """ Inner class for patterns """

components/netdriver/server/handlers/__init__.py

@@ -2,19 +2,24 @@
 # -*- coding: utf-8 -*-
 from asyncssh import SSHServerProcess
 
+from netdriver.server.handlers.arista.arista_eos import AristaEOSHandler
 from netdriver.server.handlers.array.array_ag import ArrayAGHandler
 from netdriver.server.handlers.chaitin.chaitin_ctdsg import ChaiTinCTDSGHandler
+from netdriver.server.handlers.check_point.check_point_gaia import CheckPointGaiaHandler
 from netdriver.server.handlers.cisco.cisco_asa import CiscoASAHandler
 from netdriver.server.handlers.cisco.cisco_nexus import CiscoNexusHandler
 from netdriver.server.handlers.command_handler import CommandHandler
 from netdriver.server.handlers.dptech.dptech_fw1000 import DptechFW1000Handler
+from netdriver.server.handlers.fortinet.fortinet_fortigate import FortinetFortigateHandler
+from netdriver.server.handlers.huawei.huawei_ce import HuaweiCEHandler
 from netdriver.server.handlers.huawei.huawei_usg import HuaweiUSGHandler
 from netdriver.server.handlers.hillstone.hillstone_sg6000 import HillstoneSG6000Handler
 from netdriver.server.handlers.h3c.h3c_secpath import H3CSecPathHandler
 from netdriver.server.handlers.h3c.h3c_vsr import H3CVsrHandler
 from netdriver.server.handlers.juniper.juniper_srx import JuniperSRXHandler
 from netdriver.server.handlers.juniper.juniper_ex import JuniperEXHandler
 from netdriver.server.handlers.maipu.maipu_nss import MaiPuNSSHandler
+from netdriver.server.handlers.paloalto.paloalto_pa import PaloaltoPAHandler
 from netdriver.server.handlers.qianxin.qianxin_nsg import QiAnXinNSGHandler
 from netdriver.server.handlers.topsec.topsec_ngfw import TopSecNGFWHandler
 from netdriver.server.handlers.venustech.venustech_usg import VenustechUSGHandler
@@ -56,5 +61,15 @@ def create_handler(process: SSHServerProcess,
             return ChaiTinCTDSGHandler(process, conf_path)
         elif TopSecNGFWHandler.is_selectable(vendor, model, version):
             return TopSecNGFWHandler(process, conf_path)
+        elif AristaEOSHandler.is_selectable(vendor, model, version):
+            return AristaEOSHandler(process, conf_path)
+        elif CheckPointGaiaHandler.is_selectable(vendor, model, version):
+            return CheckPointGaiaHandler(process, conf_path)
+        elif FortinetFortigateHandler.is_selectable(vendor, model, version):
+            return FortinetFortigateHandler(process, conf_path)
+        elif PaloaltoPAHandler.is_selectable(vendor, model, version):
+            return PaloaltoPAHandler(process, conf_path)
+        elif HuaweiCEHandler.is_selectable(vendor, model, version):
+            return HuaweiCEHandler(process, conf_path)
         else:
             raise ValueError(f"Unsupported device: {vendor}:{model}:{version}")

components/netdriver/server/handlers/arista/__init__.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3.10.6
+# -*- coding: utf-8 -*-
+
+from pathlib import Path
+from asyncssh import SSHServerProcess
+from netdriver.client.mode import Mode
+from netdriver.exception.server import ClientExit
+from netdriver.server.handlers.command_handler import CommandHandler
+from netdriver.server.models import DeviceBaseInfo
+
+
+class AristaEOSHandler(CommandHandler):
+    """ Arista EOS Command Handler """
+
+    info = DeviceBaseInfo(
+        vendor="arista",
+        model="eos",
+        version="*",
+        description="Arista EOS Command Handler"
+    )
+
+    @classmethod
+    def is_selectable(cls, vendor: str, model: str, version: str) -> bool:
+        # only check vendor and model, check version in the future
+        if cls.info.vendor == vendor and cls.info.model == model:
+            return True
+
+    def __init__(self, process: SSHServerProcess, conf_path: str = None):
+        # current file path
+        if conf_path is None:
+            cwd_path = Path(__file__).parent
+            conf_path = f"{cwd_path}/arista_eos.yml"
+        self.conf_path = conf_path
+        super().__init__(process)
+
+    async def switch_vsys(self, command: str) -> bool:
+        return False
+
+    async def switch_mode(self, command: str) -> bool:
+        if command not in self.config.modes[self._mode].switch_mode_cmds:
+            return False
+
+        match self._mode:
+            case Mode.LOGIN:
+                if command == "exit":
+                    # logout
+                    raise ClientExit
+                if command == "enable":
+                    self._mode = Mode.ENABLE
+                    return True
+            case Mode.ENABLE:
+                if command == "exit":
+                    # logout
+                    raise ClientExit
+                elif command == "configure terminal":
+                    # switch to config mode
+                    self._mode = Mode.CONFIG
+                    return True
+                elif command == "disable":
+                    self._mode = Mode.LOGIN
+                    return True
+            case Mode.CONFIG:
+                if command == "exit" or command == "end":
+                    # exit config mode
+                    self._mode = Mode.ENABLE
+                    return True
+                elif command == "disable":
+                    # login mode
+                    self._mode = Mode.LOGIN
+                    return True
+            case _:
+                return False
+        return False