refactor: Address code review feedback (docstrings, timeouts, complexity)

Author: hyaku0121

cortex/health/checks/disk.py

@@ -2,32 +2,38 @@
 from ..monitor import HealthCheck, CheckResult
 
 class DiskCheck(HealthCheck):
+    """Check root filesystem disk usage."""
+
     def run(self) -> CheckResult:
-        total, used, free = shutil.disk_usage("/")
-        # Calculate usage percentage
+        """
+        Calculate disk usage percentage.
+        
+        Returns:
+            CheckResult based on usage thresholds.
+        """
+        # Use _ for unused variable (free space)
+        total, used, _ = shutil.disk_usage("/")
         usage_percent = (used / total) * 100
 
         score = 100
         status = "OK"
-        details = f"{usage_percent:.1f}% Used"
         rec = None
-
-        # Scoring logic (Spec compliant)
+        
         if usage_percent > 90:
             score = 0
             status = "CRITICAL"
-            rec = "Clean package cache (+50 pts)"
+            rec = "Clean up disk space immediately"
         elif usage_percent > 80:
             score = 50
             status = "WARNING"
-            rec = "Clean package cache (+10 pts)"
-        
+            rec = "Consider cleaning up disk space"
+            
         return CheckResult(
-            name="Disk Space",
+            name="Disk Usage",
             category="disk",
             score=score,
             status=status,
-            details=details,
+            details=f"{usage_percent:.1f}% used",
             recommendation=rec,
-            weight=0.15  # 15%
+            weight=0.20
         )

cortex/health/checks/security.py

@@ -3,53 +3,32 @@
 from ..monitor import HealthCheck, CheckResult
 
 class SecurityCheck(HealthCheck):
+    """Check security configuration including firewall and SSH settings."""
+
     def run(self) -> CheckResult:
+        """
+        Run security checks for firewall status and SSH configuration.
+        
+        Returns:
+            CheckResult with security score based on detected issues.
+        """
         score = 100
         issues = []
         recommendations = []
 
         # 1. Firewall (UFW) Check
-        ufw_active = False
-        try:
-            # Add timeout to prevent hanging (Fixes Reliability Issue)
-            res = subprocess.run(
-                ["systemctl", "is-active", "ufw"], 
-                capture_output=True, 
-                text=True,
-                timeout=5
-            )
-            # Fix: Use exact match to avoid matching "inactive" which contains "active"
-            if res.returncode == 0 and res.stdout.strip() == "active":
-                ufw_active = True
-        except subprocess.TimeoutExpired:
-            pass # Command timed out, treat as inactive or unavailable
-        except FileNotFoundError:
-            pass # Environment without systemctl (e.g., Docker or non-systemd)
-        except Exception:
-            pass # Generic error protection
-
+        ufw_active, ufw_issue, ufw_rec = self._check_firewall()
         if not ufw_active:
-            score = 0 # Spec: 0 points if Firewall is inactive
-            issues.append("Firewall Inactive")
-            recommendations.append("Enable UFW Firewall")
+            score = 0
+            issues.append(ufw_issue)
+            recommendations.append(ufw_rec)
 
         # 2. SSH Root Login Check
-        try:
-            ssh_config = "/etc/ssh/sshd_config"
-            if os.path.exists(ssh_config):
-                with open(ssh_config, 'r') as f:
-                    for line in f:
-                        line = line.strip()
-                        # Check for uncommented PermitRootLogin yes
-                        if line.startswith("PermitRootLogin") and "yes" in line.split():
-                            score -= 50
-                            issues.append("Root SSH Allowed")
-                            recommendations.append("Disable SSH Root Login in sshd_config")
-                            break
-        except PermissionError:
-            pass # Cannot read config, skip check
-        except Exception:
-            pass # Generic error protection
+        ssh_penalty, ssh_issue, ssh_rec = self._check_ssh_root_login()
+        if ssh_penalty > 0:
+            score -= ssh_penalty
+            issues.append(ssh_issue)
+            recommendations.append(ssh_rec)
 
         status = "OK"
         if score < 50: status = "CRITICAL"
@@ -63,4 +42,35 @@ def run(self) -> CheckResult:
             details=", ".join(issues) if issues else "Secure",
             recommendation=", ".join(recommendations) if recommendations else None,
             weight=0.35
-        )
+        )
+
+    def _check_firewall(self):
+        """Check if UFW is active."""
+        try:
+            res = subprocess.run(
+                ["systemctl", "is-active", "ufw"], 
+                capture_output=True, 
+                text=True,
+                timeout=10
+            )
+            if res.returncode == 0 and res.stdout.strip() == "active":
+                return True, None, None
+        except (subprocess.TimeoutExpired, FileNotFoundError, Exception):
+            pass
+        
+        return False, "Firewall Inactive", "Enable UFW Firewall"
+
+    def _check_ssh_root_login(self):
+        """Check for PermitRootLogin yes in sshd_config."""
+        try:
+            ssh_config = "/etc/ssh/sshd_config"
+            if os.path.exists(ssh_config):
+                with open(ssh_config, 'r') as f:
+                    for line in f:
+                        line = line.strip()
+                        if line.startswith("PermitRootLogin") and "yes" in line.split():
+                            return 50, "Root SSH Allowed", "Disable SSH Root Login in sshd_config"
+        except (PermissionError, Exception):
+            pass
+        
+        return 0, None, None

cortex/health/checks/updates.py

@@ -2,55 +2,58 @@
 from ..monitor import HealthCheck, CheckResult
 
 class UpdateCheck(HealthCheck):
+    """Check for pending system updates and security patches."""
+
     def run(self) -> CheckResult:
+        """
+        Check for available updates using apt.
+        
+        Returns:
+            CheckResult with score based on pending updates.
+        """
         score = 100
         pkg_count = 0
         sec_count = 0
-        rec = None
 
-        # Parse apt list --upgradable
         try:
-            # Execute safely without pipeline
+            # Add timeout to prevent hangs
             res = subprocess.run(
                 ["apt", "list", "--upgradable"], 
-                capture_output=True, text=True
+                capture_output=True, 
+                text=True,
+                timeout=30
             )
-            
             lines = res.stdout.splitlines()
-            # Skip first line "Listing..."
+            
+            # apt list output header usually takes first line
             for line in lines[1:]:
                 if line.strip():
-                    pkg_count += 1
                     if "security" in line.lower():
                         sec_count += 1
+                    else:
+                        pkg_count += 1
 
             # Scoring
-            score -= (pkg_count * 2) # -2 pts per normal package
-            score -= (sec_count * 10) # -10 pts per security package
-            
-            if pkg_count > 0:
-                rec = f"Install {pkg_count} updates (+{100-score} pts)"
+            score -= (pkg_count * 2) 
+            score -= (sec_count * 10) 
 
-        except FileNotFoundError:
-            # Skip on non-apt environments (100 pts)
-            return CheckResult("Updates", "updates", 100, "SKIP", "apt not found", weight=0.30)
         except Exception:
-            pass # Ignore errors
+             pass
 
         status = "OK"
-        if score < 60: status = "CRITICAL"
-        elif score < 100: status = "WARNING"
+        if score < 50: status = "CRITICAL"
+        elif score < 90: status = "WARNING"
 
-        details = f"{pkg_count} pending"
-        if sec_count > 0:
-            details += f" ({sec_count} security)"
+        details = f"{pkg_count} packages, {sec_count} security updates pending"
+        if pkg_count == 0 and sec_count == 0:
+            details = "System up to date"
 
         return CheckResult(
             name="System Updates",
             category="updates",
             score=max(0, score),
             status=status,
             details=details,
-            recommendation=rec,
-            weight=0.30 # 30%
+            recommendation="Run 'apt upgrade'" if score < 100 else None,
+            weight=0.25
         )