From 3b3a915552de45e93f0b593134c8301d92280a6c Mon Sep 17 00:00:00 2001
From: He Yin Ling <heyinling@espressif.com>
Date: Sun, 7 Jan 2018 20:24:38 +0800
Subject: [PATCH] tiny-test-fw: implement `env.get_pc_nic_info`:

this method is used to get mac/ipv4/ipv6 address for one NIC
---
 tools/tiny-test-fw/Env.py         | 31 +++++++++++++++++++++++++------
 tools/tiny-test-fw/docs/index.rst |  1 +
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/tools/tiny-test-fw/Env.py b/tools/tiny-test-fw/Env.py
index da5c0b598..fa194d10c 100644
--- a/tools/tiny-test-fw/Env.py
+++ b/tools/tiny-test-fw/Env.py
@@ -17,6 +17,8 @@ import os
 import threading
 import functools
 
+import netifaces
+
 import EnvConfig
 
 
@@ -130,17 +132,34 @@ class Env(object):
         """
         return self.config.get_variable(variable_name)
 
+    PROTO_MAP = {
+        "ipv4": netifaces.AF_INET,
+        "ipv6": netifaces.AF_INET6,
+        "mac": netifaces.AF_LINK,
+    }
+
     @_synced
-    def get_pc_nic_info(self, nic_name="pc_nic"):
+    def get_pc_nic_info(self, nic_name="pc_nic", proto="ipv4"):
         """
         get_pc_nic_info(nic_name="pc_nic")
-        try to get nic info (ip address, ipv6 address, mac address)
+        try to get info of a specified NIC and protocol.
 
-        :param nic_name: pc nic name. allows passing variable name, nic name value or omitted (to get default nic info).
-        :return: a dict of address ("ipv4", "ipv6", "mac") if successfully found. otherwise None.
+        :param nic_name: pc nic name. allows passing variable name, nic name value.
+        :param proto: "ipv4", "ipv6" or "mac"
+        :return: a dict of nic info if successfully found. otherwise None.
+                 nic info keys could be different for different protocols.
+                 key "addr" is available for both mac, ipv4 and ipv6 pic info.
         """
-        # TODO: need to implement auto get nic info method
-        return self.config.get_variable("nic_info/" + nic_name)
+        interfaces = netifaces.interfaces()
+        if nic_name in interfaces:
+            # the name is in the interface list, we regard it as NIC name
+            if_addr = netifaces.ifaddresses(nic_name)
+        else:
+            # it's not in interface name list, we assume it's variable name
+            _nic_name = self.get_variable(nic_name)
+            if_addr = netifaces.ifaddresses(_nic_name)
+
+        return if_addr[self.PROTO_MAP[proto]][0]
 
     @_synced
     def close(self):
diff --git a/tools/tiny-test-fw/docs/index.rst b/tools/tiny-test-fw/docs/index.rst
index a70e6bf32..22cfd73a8 100644
--- a/tools/tiny-test-fw/docs/index.rst
+++ b/tools/tiny-test-fw/docs/index.rst
@@ -128,6 +128,7 @@ The following 3rd party lib is required:
     * pyserial
     * pyyaml
     * xunitgen
+    * netifaces
 
 To build document, we need to install ``Sphinx`` and ``sphinx-rtd-theme`` (you may replace this with your own theme).