From bd9a848737bf9dada185ff5579f64658bc2bfd75 Mon Sep 17 00:00:00 2001
From: Swapnil <swapnil.theskss@gmail.com>
Date: Sun, 23 Jun 2024 14:26:58 +0530
Subject: [PATCH] updated cpu_count - use dmedecode

---
 agent.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/agent.py b/agent.py
index f21703e..4f0e9fc 100644
--- a/agent.py
+++ b/agent.py
@@ -38,9 +38,20 @@ class ServerData:
         return ram, disk
 
     def get_cpu_count(self):
-        with open('/proc/cpuinfo', 'r') as f:
-            cpuinfo = f.read()
-        cpu_count = cpuinfo.count('processor')
+        cpu_count = 0
+        if os.path.isfile('/usr/sbin/dmidecode'):
+            try:
+                output = subprocess.check_output(['sudo', '/usr/sbin/dmidecode', '-t', 'processor']).decode('utf-8')
+                core_match = re.search(r'Core Count: (\d+)', output)
+                thread_match = re.search(r'Thread Count: (\d+)', output)
+                if core_match and thread_match:
+                    cpu_count = int(core_match.group(1)) * int(thread_match.group(1))
+            except subprocess.CalledProcessError:
+                pass
+        if cpu_count == 0:
+            with open('/proc/cpuinfo', 'r') as f:
+                cpuinfo = f.read()
+            cpu_count = cpuinfo.count('processor')
         logging.info(f"CPU Count: {cpu_count}")
         return cpu_count