From 7b6b2b9150e7970492594f078ea957b7826135f1 Mon Sep 17 00:00:00 2001 From: CyberRoute Date: Mon, 23 Sep 2024 14:23:11 +0200 Subject: [PATCH] refactoring vendor.py --- core/vendor.py | 85 +++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 78 insertions(+), 7 deletions(-) diff --git a/core/vendor.py b/core/vendor.py index cee7da6..0ced21d 100644 --- a/core/vendor.py +++ b/core/vendor.py @@ -1,3 +1,6 @@ +""" +The MacVendorLookup load data about vendors from an oui.csv file +""" import os import csv from io import StringIO @@ -5,10 +8,29 @@ 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") \ No newline at end of file + return self.mac_vendor_data.get(oui, "Vendor not found")