From 49a3f63c34e5a8e28700f9851612009b230274d1 Mon Sep 17 00:00:00 2001 From: Shailaja Date: Fri, 5 Jul 2024 15:03:11 +0530 Subject: [PATCH] refractor --- agent.py | 76 +++++++++++++++++++++++++------------------------------- 1 file changed, 34 insertions(+), 42 deletions(-) diff --git a/agent.py b/agent.py index 16b3f23..26751ff 100644 --- a/agent.py +++ b/agent.py @@ -264,48 +264,44 @@ class ServerData: def get_os_id(self, os_list): - try: - os_info = self.get_os_release_info() - if not os_info: - logging.error("No OS release info found.") - for os_entry in os_list: - if os_entry['name'].lower() in ["other", "custom"]: - return os_entry['id'] - return None - - if 'NAME' in os_info and 'VERSION' in os_info: - current_os = f"{os_info.get('NAME', 'Unknown')} {os_info.get('VERSION', '').strip()}".strip().lower() - else: - current_os = f"{os_info.get('os_name', 'Unknown')} {os_info.get('productversion', '').strip()}".strip().lower() - - for os_entry in os_list: - if current_os in os_entry['name'].lower(): - return os_entry['id'] - - # Fallback checks for common OS names if full name doesn't match - for os_entry in os_list: - if 'ubuntu' in current_os and 'ubuntu' in os_entry['name'].lower(): - return os_entry['id'] - elif 'centos' in current_os and 'centos' in os_entry['name'].lower(): - return os_entry['id'] - elif 'fedora' in current_os and 'fedora' in os_entry['name'].lower(): - return os_entry['id'] - - # Default to 'other' or 'custom' if no match found + os_info = self.get_os_release_info() + + if not os_info: + logging.error("No OS release info found.") for os_entry in os_list: if os_entry['name'].lower() in ["other", "custom"]: return os_entry['id'] - return None - except Exception as e: - logging.error("Failed to fetch OS ID: {}".format(e)) - for os_entry in os_list: - if os_entry['name'].lower() in ["other", "custom"]: - return os_entry['id'] - return None + return 1 + + if 'NAME' in os_info and 'VERSION' in os_info: + current_os = f"{os_info.get('NAME', 'Unknown')} {os_info.get('VERSION', '').strip()}".strip().lower() + else: + current_os = f"{os_info.get('os_name', 'Unknown')} {os_info.get('productversion', '').strip()}".strip().lower() + + for os_entry in os_list: + if current_os in os_entry['name'].lower(): + return os_entry['id'] + + # Fallback checks for common OS names if full name doesn't match + for os_entry in os_list: + if 'ubuntu' in current_os and 'ubuntu' in os_entry['name'].lower(): + return os_entry['id'] + elif 'centos' in current_os and 'centos' in os_entry['name'].lower(): + return os_entry['id'] + elif 'fedora' in current_os and 'fedora' in os_entry['name'].lower(): + return os_entry['id'] + + # Default to 'other' or 'custom' if no match found + for os_entry in os_list: + if os_entry['name'].lower() in ["other", "custom"]: + return os_entry['id'] + return 1 + def create_post_data(self): ram, disk = self.get_ram_and_disk() + os_list = self.get_os_list() post_data = { "server_type": 1, "os_id": self.get_os_id(os_list), @@ -471,14 +467,10 @@ class ServerManager: else: return self.create_note(note_data) - def get_os_list(self): + def get_os_list(self): os_list = self.send_request('GET', '/api/v1/os') - if os_list: - logging.info("OS list fetched successfully") - return os_list - else: - logging.error("Failed to fetch OS list") - return [] + logging.info("OS list fetched successfully") if os_list else logging.error("Failed to fetch OS list") + return os_list or [] def validate_env_vars(): api_key = os.getenv('AGENT_API')