refactoring vendor.py

core/vendor.py

@@ -1,14 +1,36 @@
+"""
+The MacVendorLookup load data about vendors from an oui.csv file
+"""
 import os
 import csv
 from io import StringIO
 import requests
 
 
 class MacVendorLookup:
+    """Mac2Vendor class for translation of a MAC address to its vendor.
+
+    This class looks up the vendor of a given MAC address using the
+    Organizationally Unique Identifier (OUI) part of the MAC address.
+    Data can be loaded either from a URL or a cached CSV file.
+
+    Attributes:
+        mac_vendor_data (dict): A dictionary mapping OUI prefixes to vendor names.
+    """
+
     mac_vendor_data = None
 
     @classmethod
     def load_data(cls, url):
+        """Loads MAC address to vendor data, either from a cached file or a URL.
+
+        If the data is already loaded, this method does nothing.
+        If a cache file (oui.csv) exists, the data is loaded from the file.
+        Otherwise, the data is fetched from the provided URL and saved to the cache.
+
+        Args:
+            url (str): The URL from which to download the MAC vendor data.
+        """
         if cls.mac_vendor_data is None:
             cache_file = "oui.csv"
             if os.path.exists(cache_file):
@@ -19,7 +41,17 @@ def load_data(cls, url):
 
     @classmethod
     def _load_from_file(cls, filename):
-        with open(filename, 'r') as file:
+        """Loads MAC vendor data from a local CSV file.
+
+        The CSV file should have a header row followed by rows of OUI and vendor information.
+
+        Args:
+            filename (str): The path to the CSV file to read from.
+
+        Returns:
+            dict: A dictionary mapping OUI (first 6 characters of MAC address) to vendor names.
+        """
+        with open(filename, 'r', encoding='utf-8') as file:
             csvreader = csv.reader(file)
             next(csvreader)  # Skip header
             mac_vendor_data = {}
@@ -31,7 +63,19 @@ def _load_from_file(cls, filename):
 
     @classmethod
     def _load_from_url(cls, url):
-        response = requests.get(url)
+        """Loads MAC vendor data from a given URL.
+
+        Fetches the CSV data from the provided URL and parses it into
+        a dictionary of OUI to vendor mappings.
+
+        Args:
+            url (str): The URL to fetch the MAC vendor data from.
+
+        Returns:
+            dict: A dictionary mapping OUI (first 6 characters of MAC address) to vendor names.
+            If the request fails, returns an empty dictionary.
+        """
+        response = requests.get(url, timeout=5)
         if response.status_code == 200:
             csv_data = StringIO(response.text)
             csvreader = csv.reader(csv_data)
@@ -42,22 +86,49 @@ def _load_from_url(cls, url):
                 vendor = row[2]
                 mac_vendor_data[oui] = vendor
             return mac_vendor_data
-        else:
-            print(f"Failed to fetch data from {url}. Status code: {response.status_code}")
-            return {}
+        print(f"Failed to fetch data from {url}. Status code: {response.status_code}")
+        return {}
 
     @classmethod
     def _save_to_file(cls, filename, data):
-        with open(filename, 'w') as file:
+        """Saves the MAC vendor data to a local CSV file.
+
+        The file will contain two columns: OUI and vendor name. This method
+        will overwrite the file if it already exists.
+
+        Args:
+            filename (str): The path to the file where the data should be saved.
+            data (dict): A dictionary containing OUI to vendor mappings.
+        """
+        with open(filename, 'w', encoding='utf-8') as file:
             writer = csv.writer(file)
             writer.writerow(["Organizationally Unique Identifier", "Organization Name"])
             for oui, vendor in data.items():
                 writer.writerow([oui, vendor])
 
     def __init__(self, url):
+        """Initializes the MacVendorLookup class by loading the MAC vendor data.
+
+        Args:
+            url (str): The URL from which to fetch the MAC vendor data if the cache doesn't exist.
+        """
         self.load_data(url)
 
     def lookup_vendor(self, mac_address):
+        """Looks up the vendor for a given MAC address.
+
+        The method extracts the first six characters (OUI) from the MAC address,
+        which identifies the vendor. If the OUI is found in the loaded data, the
+        corresponding vendor is returned; otherwise, it returns "Vendor not found".
+
+        Args:
+            mac_address (str): The MAC address to look up, in formats like
+            XX:XX:XX:YY:YY:YY or XX-XX-XX-YY-YY-YY.
+
+        Returns:
+            str: The vendor name associated with the OUI, or "Vendor not found"
+            if the OUI is not recognized.
+        """
         cleaned_mac = mac_address.upper().replace(":", "").replace("-", "")
         oui = cleaned_mac[:6]
-        return self.mac_vendor_data.get(oui, "Vendor not found")
+        return self.mac_vendor_data.get(oui, "Vendor not found")