[dts] [‘dts-v1’ 2/9] Optimize ssh connection

[sjiajiax <sunx.jiajia@intel.com>]

Optimize configure parse

Move some functions in dts to an common utils module

Signed-off-by: sjiajiax <sunx.jiajia@intel.com>
---
 framework/config.py         | 171 +++++++++++++++++++++++++++++++++++---------
 framework/ssh_connection.py |   5 ++
 framework/ssh_pexpect.py    |  63 +++++++++++++---
 framework/utils.py          |  77 ++++++++++++++++++++
 4 files changed, 270 insertions(+), 46 deletions(-)
 mode change 100755 => 100644 framework/config.py
 create mode 100644 framework/utils.py

diff --git a/framework/config.py b/framework/config.py
old mode 100755
new mode 100644
index d2548e8..927c846
--- a/framework/config.py
+++ b/framework/config.py
@@ -37,45 +37,133 @@ import re
 import ConfigParser  # config parse module
 import argparse      # prase arguments module
 
-portconf = "conf/ports.cfg"
-crbconf = "conf/crbs.cfg"
+PORTCONF = "conf/ports.cfg"
+CRBCONF = "conf/crbs.cfg"
+VIRTCONF = "conf/virt_global.cfg"
 
 
 class UserConf():
 
-    def __init__(self, port_conf=portconf, crb_conf=crbconf):
-        self.port_config = port_conf
-        self.crb_config = crb_conf
+    def __init__(self, config):
+        self.conf = ConfigParser.SafeConfigParser()
+        load_files = self.conf.read(config)
+        if load_files == []:
+            print "FAILED LOADING %s!!!" % config
+            self.conf = None
+            raise
+
+    def get_sections(self):
+        if self.conf is None:
+            return None
+
+        return self.conf.sections()
+
+    def load_section(self, section):
+        if self.conf is None:
+            return None
+
+        items = None
+        for conf_sect in self.conf.sections():
+            if conf_sect == section:
+                items = self.conf.items(section)
+
+        return items
+
+    def load_config(self, item):
+        confs = [conf.strip() for conf in item.split(';')]
+        if '' in confs:
+            confs.remove('')
+        return confs
+
+    def load_param(self, conf):
+        paramDict = dict()
+
+        for param in conf.split(','):
+            (key, _, value) = param.partition('=')
+            paramDict[key] = value
+        return paramDict
+
+
+class VirtConf(UserConf):
+
+    def __init__(self, virt_conf=VIRTCONF):
+        self.config_file = virt_conf
+        self.virt_cfg = {}
+        try:
+            self.virt_conf = UserConf(self.config_file)
+        except Exception as e:
+            print "FAILED LOADING VIRT CONFIG!!!"
+            self.virt_conf = None
+
+    def load_virt_config(self, name):
+        self.virt_cfgs = []
+
+        try:
+            virt_confs = self.virt_conf.load_section(name)
+        except:
+            print "FAILED FIND SECTION %s!!!" % name
+            return
+
+        for virt_conf in virt_confs:
+            virt_cfg = {}
+            virt_params = []
+            key, config = virt_conf
+            confs = self.virt_conf.load_config(config)
+            for config in confs:
+                virt_params.append(self.load_virt_param(config))
+            virt_cfg[key] = virt_params
+            self.virt_cfgs.append(virt_cfg)
+
+    def get_virt_config(self):
+        return self.virt_cfgs
+
+    def load_virt_param(self, config):
+        cfg_params = self.virt_conf.load_param(config)
+        return cfg_params
+
+
+class PortConf(UserConf):
+
+    def __init__(self, port_conf=PORTCONF):
+        self.config_file = port_conf
         self.ports_cfg = {}
         self.pci_regex = "([\da-f]{2}:[\da-f]{2}.\d{1})$"
         try:
-            self.port_conf = ConfigParser.SafeConfigParser()
-            self.port_conf.read(self.port_config)
+            self.port_conf = UserConf(self.config_file)
         except Exception as e:
             print "FAILED LOADING PORT CONFIG!!!"
+            self.port_conf = None
 
     def load_ports_config(self, crbIP):
-        ports = []
-        for crb in self.port_conf.sections():
-            if crb != crbIP:
-                continue
-            ports = [port.strip()
-                     for port in self.port_conf.get(crb, 'ports').split(';')]
+        self.ports_cfg = {}
+        if self.port_conf is None:
+            return
+
+        ports = self.port_conf.load_section(crbIP)
+        if ports is None:
+            return
+        key, config = ports[0]
+        confs = self.port_conf.load_config(config)
+
+        for config in confs:
+            port_param = self.port_conf.load_param(config)
 
-        for port in ports:
-            port_cfg = self.__parse_port_param(port)
             # check pci BDF validity
-            if 'pci' not in port_cfg:
+            if 'pci' not in port_param:
                 print "NOT FOUND CONFIG FOR NO PCI ADDRESS!!!"
                 continue
-            m = re.match(self.pci_regex, port_cfg['pci'])
+            m = re.match(self.pci_regex, port_param['pci'])
             if m is None:
                 print "INVALID CONFIG FOR NO PCI ADDRESS!!!"
                 continue
 
-            keys = port_cfg.keys()
+            keys = port_param.keys()
             keys.remove('pci')
-            self.ports_cfg[port_cfg['pci']] = {key: port_cfg[key] for key in keys}
+            self.ports_cfg[port_param['pci']] = {
+                key: port_param[key] for key in keys}
+            if 'numa' in self.ports_cfg[port_param['pci']]:
+                numa_str = self.ports_cfg[port_param['pci']]['numa']
+                self.ports_cfg[port_param['pci']]['numa'] = int(numa_str)
 
     def get_ports_config(self):
         return self.ports_cfg
@@ -86,23 +174,36 @@ class UserConf():
         else:
             return False
 
-    def __parse_port_param(self, port):
-        portDict = dict()
-
-        for param in port.split(','):
-            (key, _, value) = param.partition('=')
-            if key == 'numa':
-                portDict[key] = int(value)
-            else:
-                portDict[key] = value
-        return portDict
+  
 
 
 if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description="Load DTS configuration files")
-    parser.add_argument("-p", "--portconf", default=portconf)
-    parser.add_argument("-c", "--crbconf", default=crbconf)
+    parser = argparse.ArgumentParser(
+        description="Load DTS configuration files")
+    parser.add_argument("-p", "--portconf", default=PORTCONF)
+    parser.add_argument("-c", "--crbconf", default=CRBCONF)
+    parser.add_argument("-v", "--virtconf", default=VIRTCONF)
     args = parser.parse_args()
-    conf = UserConf()
-    conf.load_ports_config('192.168.1.1')
-    conf.check_port_available('0000:86:00.0')
+
+    # not existed configuration file
+    VirtConf('/tmp/not-existed.cfg')
+
+    # example for basic use configuration file
+    conf = UserConf(PORTCONF)
+    for section in conf.get_sections():
+        items = conf.load_section(section)
+        key, value = items[0]
+        confs = conf.load_config(value)
+        for config in confs:
+            conf.load_param(config)
+
+    # example for port configuration file
+    portconf = PortConf(PORTCONF)
+    portconf.load_ports_config('DUT IP')
+    print portconf.get_ports_config()
+    portconf.check_port_available('86:00.0')
+
+    # example for global virtualization configuration file
+    virtconf = VirtConf(VIRTCONF)
+    virtconf.load_virt_config('LIBVIRT')
+    print virtconf.get_virt_config()
diff --git a/framework/ssh_connection.py b/framework/ssh_connection.py
index 18a6517..7286b14 100644
--- a/framework/ssh_connection.py
+++ b/framework/ssh_connection.py
@@ -62,6 +62,11 @@ class SSHConnection(object):
         self.logger.debug(out)
         return out
 
+    def get_session_before(self, timeout=15):
+        out = self.session.get_session_before(timeout)
+        self.logger.debug(out)
+        return out
+
     def close(self):
         self.session.close()
 
diff --git a/framework/ssh_pexpect.py b/framework/ssh_pexpect.py
index 735df44..4193020 100644
--- a/framework/ssh_pexpect.py
+++ b/framework/ssh_pexpect.py
@@ -2,7 +2,8 @@ import time
 import pexpect
 import pxssh
 from debugger import ignore_keyintr, aware_keyintr
-from exception import TimeoutException, SSHConnectionException
+from exception import TimeoutException, SSHConnectionException, SSHSessionDeadException
+from utils import RED, GREEN
 
 """
 Module handle ssh sessions between tester and DUT.
@@ -14,16 +15,28 @@ Aslo support transfer files to tester or DUT.
 class SSHPexpect(object):
 
     def __init__(self, host, username, password):
-        self.magic_prompt = "[MAGIC PROMPT]"
+        self.magic_prompt = "MAGIC PROMPT"
         try:
             self.session = pxssh.pxssh()
-            self.username = username
             self.host = host
+            self.username = username
             self.password = password
-            self.session.login(self.host, self.username,
-                               self.password, original_prompt='[$#>]')
+            if ':' in host:
+                self.ip = host.split(':')[0]
+                self.port = int(host.split(':')[1])
+                self.session.login(self.ip, self.username,
+                                   self.password, original_prompt='[$#>]',
+                                   port=self.port, login_timeout=20)
+            else:
+                self.session.login(self.host, self.username,
+                                   self.password, original_prompt='[$#>]')
             self.send_expect('stty -echo', '# ', timeout=2)
-        except Exception:
+        except Exception, e:
+            print RED(e)
+            if getattr(self, 'port', None):
+                suggestion = "\nSuggession: Check if the fireware on [ %s ] " % \
+                    self.ip + "is stoped\n"
+                print GREEN(suggestion)
             raise SSHConnectionException(host)
 
     def init_log(self, logger, name):
@@ -33,7 +46,7 @@ class SSHPexpect(object):
 
     def send_expect_base(self, command, expected, timeout=15):
         ignore_keyintr()
-        self.__flush() # clear buffer
+        self.__flush()  # clear buffer
         self.session.PROMPT = expected
         self.__sendline(command)
         self.__prompt(command, timeout)
@@ -45,7 +58,7 @@ class SSHPexpect(object):
     def send_expect(self, command, expected, timeout=15, verify=False):
         ret = self.send_expect_base(command, expected, timeout)
         if verify:
-            ret_status = self.send_expect_base("echo $?", expected)
+            ret_status = self.send_expect_base("echo $?", expected, timeout)
             if not int(ret_status):
                 return ret
             else:
@@ -54,21 +67,44 @@ class SSHPexpect(object):
         else:
             return ret
 
-    def __flush(self):
+    def get_session_before(self, timeout=15):
+        """
+        Get all output before timeout
+        """
+        ignore_keyintr()
         self.session.PROMPT = self.magic_prompt
-        self.session.prompt(0.1)
+        try:
+            self.session.prompt(timeout)
+        except Exception as e:
+            pass
+
+        aware_keyintr()
+        before = self.get_output_before()
+        self.__flush()
+        return before
+
+    def __flush(self):
+        """
+        Clear all session buffer
+        """
+        self.session.buffer = ""
+        self.session.before = ""
 
     def __prompt(self, command, timeout):
         if not self.session.prompt(timeout):
             raise TimeoutException(command, self.get_output_all())
 
     def __sendline(self, command):
+        if not self.isalive():
+            raise SSHSessionDeadException(self.host)
         if len(command) == 2 and command.startswith('^'):
             self.session.sendcontrol(command[1])
         else:
             self.session.sendline(command)
 
     def get_output_before(self):
+        if not self.isalive():
+            raise SSHSessionDeadException(self.host)
         self.session.flush()
         before = self.session.before.rsplit('\r\n', 1)
         if before[0] == "[PEXPECT]":
@@ -103,7 +139,12 @@ class SSHPexpect(object):
         """
         Sends a local file to a remote place.
         """
-        command = 'scp {0} {1}@{2}:{3}'.format(src, self.username, self.host, dst)
+        if ':' in self.host:
+            command = 'scp -P {0} -o NoHostAuthenticationForLocalhost=yes {1} {2}@{3}:{4}'.format(
+                str(self.port), src, self.username, self.ip, dst)
+        else:
+            command = 'scp {0} {1}@{2}:{3}'.format(
+                src, self.username, self.host, dst)
         if password == '':
             self._spawn_scp(command, self.password)
         else:
diff --git a/framework/utils.py b/framework/utils.py
new file mode 100644
index 0000000..57eb988
--- /dev/null
+++ b/framework/utils.py
@@ -0,0 +1,77 @@
+# BSD LICENSE
+#
+# Copyright(c) 2010-2014 Intel Corporation. All rights reserved.
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#
+#   * Redistributions of source code must retain the above copyright
+#     notice, this list of conditions and the following disclaimer.
+#   * Redistributions in binary form must reproduce the above copyright
+#     notice, this list of conditions and the following disclaimer in
+#     the documentation and/or other materials provided with the
+#     distribution.
+#   * Neither the name of Intel Corporation nor the names of its
+#     contributors may be used to endorse or promote products derived
+#     from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import json         # json format
+import re
+
+
+def RED(text):
+    return "\x1B[" + "31;1m" + str(text) + "\x1B[" + "0m"
+
+
+def BLUE(text):
+    return "\x1B[" + "36;1m" + str(text) + "\x1B[" + "0m"
+
+
+def GREEN(text):
+    return "\x1B[" + "32;1m" + str(text) + "\x1B[" + "0m"
+
+
+def pprint(some_dict):
+    """
+    Print JSON format dictionary object.
+    """
+    return json.dumps(some_dict, sort_keys=True, indent=4)
+
+
+def regexp(s, to_match, allString=False):
+    """
+    Ensure that the re `to_match' only has one group in it.
+    """
+
+    scanner = re.compile(to_match, re.DOTALL)
+    if allString:
+        return scanner.findall(s)
+    m = scanner.search(s)
+    if m is None:
+        print RED("Failed to match " + to_match + " in the string " + s)
+        return None
+    return m.group(1)
+
+
+def get_obj_funcs(obj, func_name_regex):
+    """
+    Return function list which name matched regex.
+    """
+    for func_name in dir(obj):
+        func = getattr(obj, func_name)
+        if callable(func) and re.match(func_name_regex, func.__name__):
+            yield func
-- 
1.9.0