Enhance bandwidth usage calculation with resource management

- Introduced memory and processing time limits to prevent system overload during bandwidth calculations.
- Added error handling for file operations and improved logging for better traceability.
- Implemented batch processing of log lines to manage memory usage effectively.
- Updated methods to safely parse log data and handle large files, ensuring robust performance.
- Refactored code for clarity and maintainability, including the addition of helper functions for file size and memory limit settings.

plogical/findBWUsage.py

@@ -1,65 +1,157 @@
 import sys
 sys.path.append('/usr/local/CyberCP')
 import os
+import gc
+import time
 from plogical import CyberCPLogFileWriter as logging
 import shlex
 import subprocess
 import validators
+import resource
 
 class findBWUsage:
+    # Configuration constants
+    MAX_MEMORY_MB = 512  # Maximum memory usage in MB
+    MAX_PROCESSING_TIME = 300  # Maximum processing time in seconds (5 minutes)
+    MAX_LOG_LINES_PER_BATCH = 10000  # Process logs in batches
+    MAX_FILE_SIZE_MB = 100  # Skip files larger than 100MB
+    
     @staticmethod
     def parse_last_digits(line):
-        return line.split(' ')
+        """Safely parse log line and extract bandwidth data"""
+        try:
+            parts = line.split(' ')
+            if len(parts) < 10:
+                return None
+            # Extract the size field (index 9) and clean it
+            size_str = parts[9].replace('"', '').strip()
+            if size_str == '-':
+                return 0
+            return int(size_str)
+        except (ValueError, IndexError, AttributeError):
+            return None
+
+    @staticmethod
+    def get_file_size_mb(filepath):
+        """Get file size in MB"""
+        try:
+            return os.path.getsize(filepath) / (1024 * 1024)
+        except OSError:
+            return 0
+
+    @staticmethod
+    def set_memory_limit():
+        """Set memory limit to prevent system overload"""
+        try:
+            # Set memory limit to MAX_MEMORY_MB
+            memory_limit = findBWUsage.MAX_MEMORY_MB * 1024 * 1024  # Convert to bytes
+            resource.setrlimit(resource.RLIMIT_AS, (memory_limit, memory_limit))
+        except Exception as e:
+            logging.CyberCPLogFileWriter.writeToFile(f"Failed to set memory limit: {str(e)}")
 
     @staticmethod
     def calculateBandwidth(domainName):
+        """Calculate bandwidth usage for a domain with memory protection"""
+        start_time = time.time()
+        
         try:
-            path = "/home/"+domainName+"/logs/"+domainName+".access_log"
+            path = "/home/" + domainName + "/logs/" + domainName + ".access_log"
             
             if not os.path.exists(path):
                 return 0
-            from processUtilities import ProcessUtilities
-            logData = ProcessUtilities.outputExecutioner('cat %s' % (path), 'nobody').splitlines()
-            logDataLines = len(logData)
                 
-            if not os.path.exists("/home/"+domainName+"/logs"):
+            # Check file size before processing
+            file_size_mb = findBWUsage.get_file_size_mb(path)
+            if file_size_mb > findBWUsage.MAX_FILE_SIZE_MB:
+                logging.CyberCPLogFileWriter.writeToFile(f"Skipping large file {path} ({file_size_mb:.2f}MB)")
+                return 0
+
+            if not os.path.exists("/home/" + domainName + "/logs"):
                 return 0
 
             bwmeta = "/home/cyberpanel/%s.bwmeta" % (domainName)
             
-            if not os.path.exists(path):
+            # Initialize metadata
-                writeMeta = open(bwmeta, 'w')
+            currentUsed = 0
-                writeMeta.writelines('0\n0\n')
+            currentLinesRead = 0
-                writeMeta.close()
-                os.chmod(bwmeta, 0o600)
-                return 1
             
+            # Read existing metadata
             if os.path.exists(bwmeta):
-                data = open(bwmeta).readlines()
+                try:
-                currentUsed = int(data[0].strip("\n"))
+                    with open(bwmeta, 'r') as f:
-                currentLinesRead = int(data[1].strip("\n"))
+                        data = f.readlines()
-                if currentLinesRead > logDataLines:
+                    if len(data) >= 2:
+                        currentUsed = int(data[0].strip("\n"))
+                        currentLinesRead = int(data[1].strip("\n"))
+                except (ValueError, IndexError):
+                    currentUsed = 0
                     currentLinesRead = 0
-            else:
-                currentUsed = 0
-                currentLinesRead = 0
 
-            startLine = currentLinesRead
+            # Process log file in streaming mode to avoid memory issues
+            try:
+                with open(path, 'r', encoding='utf-8', errors='ignore') as logfile:
+                    # Skip to the last processed line
+                    for _ in range(currentLinesRead):
+                        try:
+                            next(logfile)
+                        except StopIteration:
+                            break
                     
-            for line in logData[startLine:]:
+                    lines_processed = 0
-                line = line.strip('"\n')
+                    batch_size = 0
-                currentLinesRead = currentLinesRead + 1
-                if len(line)>10:
-                    currentUsed = int(findBWUsage.parse_last_digits(line)[9].replace('"', '')) + currentUsed
                     
-            writeMeta = open(bwmeta,'w')
+                    for line in logfile:
-            writeMeta.writelines(str(currentUsed)+"\n")
+                        # Check processing time limit
-            writeMeta.writelines(str(currentLinesRead) + "\n")
+                        if time.time() - start_time > findBWUsage.MAX_PROCESSING_TIME:
-            writeMeta.close()
+                            logging.CyberCPLogFileWriter.writeToFile(f"Processing timeout for {domainName}")
+                            break
                         
-            os.chmod(bwmeta, 0o600)
+                        line = line.strip()
+                        if len(line) > 10:
+                            bandwidth = findBWUsage.parse_last_digits(line)
+                            if bandwidth is not None:
+                                currentUsed += bandwidth
                             
-        except BaseException as msg:
+                            currentLinesRead += 1
+                            lines_processed += 1
+                            batch_size += 1
+                            
+                            # Process in batches to manage memory
+                            if batch_size >= findBWUsage.MAX_LOG_LINES_PER_BATCH:
+                                # Force garbage collection
+                                gc.collect()
+                                batch_size = 0
+                                
+                                # Check memory usage
+                                try:
+                                    import psutil
+                                    process = psutil.Process()
+                                    memory_mb = process.memory_info().rss / (1024 * 1024)
+                                    if memory_mb > findBWUsage.MAX_MEMORY_MB:
+                                        logging.CyberCPLogFileWriter.writeToFile(f"Memory limit reached for {domainName}")
+                                        break
+                                except ImportError:
+                                    pass  # psutil not available, continue processing
+                                
+            except (IOError, OSError) as e:
+                logging.CyberCPLogFileWriter.writeToFile(f"Error reading log file {path}: {str(e)}")
+                return 0
+
+            # Write updated metadata
+            try:
+                with open(bwmeta, 'w') as f:
+                    f.write(f"{currentUsed}\n{currentLinesRead}\n")
+                os.chmod(bwmeta, 0o600)
+            except (IOError, OSError) as e:
+                logging.CyberCPLogFileWriter.writeToFile(f"Error writing metadata {bwmeta}: {str(e)}")
+                return 0
+
+            # Log processing statistics
+            processing_time = time.time() - start_time
+            if processing_time > 10:  # Log if processing took more than 10 seconds
+                logging.CyberCPLogFileWriter.writeToFile(f"Processed {domainName}: {lines_processed} lines in {processing_time:.2f}s")
+
+        except Exception as msg:
             logging.CyberCPLogFileWriter.writeToFile(str(msg) + " [calculateBandwidth]")
             return 0
 
@@ -67,67 +159,95 @@ class findBWUsage:
 
     @staticmethod
     def startCalculations():
+        """Start bandwidth calculations with resource protection"""
         try:
+            # Set memory limit
+            findBWUsage.set_memory_limit()
+            
+            start_time = time.time()
+            domains_processed = 0
+            
             for directories in os.listdir("/home"):
+                # Check overall processing time
+                if time.time() - start_time > findBWUsage.MAX_PROCESSING_TIME * 2:
+                    logging.CyberCPLogFileWriter.writeToFile("Overall processing timeout reached")
+                    break
+                
                 if validators.domain(directories):
-                    findBWUsage.calculateBandwidth(directories)
+                    try:
-        except BaseException as msg:
+                        result = findBWUsage.calculateBandwidth(directories)
+                        domains_processed += 1
+                        
+                        # Force garbage collection after each domain
+                        gc.collect()
+                        
+                        # Small delay to prevent system overload
+                        time.sleep(0.1)
+                        
+                    except Exception as e:
+                        logging.CyberCPLogFileWriter.writeToFile(f"Error processing domain {directories}: {str(e)}")
+                        continue
+            
+            total_time = time.time() - start_time
+            logging.CyberCPLogFileWriter.writeToFile(f"Bandwidth calculation completed: {domains_processed} domains in {total_time:.2f}s")
+            
+        except Exception as msg:
             logging.CyberCPLogFileWriter.writeToFile(str(msg) + " [startCalculations]")
             return 0
 
     @staticmethod
-    def findDomainBW(domainName,totalAllowed):
+    def findDomainBW(domainName, totalAllowed):
+        """Find domain bandwidth usage with improved error handling"""
         try:
-            path = "/home/"+domainName+"/logs/"+domainName+".access_log"
+            path = "/home/" + domainName + "/logs/" + domainName + ".access_log"
 
-            if not os.path.exists("/home/"+domainName+"/logs"):
+            if not os.path.exists("/home/" + domainName + "/logs"):
-                return [0,0]
+                return [0, 0]
 
-            bwmeta = "/home/" + domainName + "/logs/bwmeta"
+            bwmeta = "/home/cyberpanel/%s.bwmeta" % (domainName)
 
             if not os.path.exists(path):
-                return [0,0]
+                return [0, 0]
-
-
 
             if os.path.exists(bwmeta):
                 try:
-                    data = open(bwmeta).readlines()
+                    with open(bwmeta, 'r') as f:
-                    currentUsed = int(data[0].strip("\n"))
+                        data = f.readlines()
                     
-                    inMB = int(float(currentUsed)/(1024.0*1024.0))
+                    if len(data) < 1:
+                        return [0, 0]
+                        
+                    currentUsed = int(data[0].strip("\n"))
+                    inMB = int(float(currentUsed) / (1024.0 * 1024.0))
+
+                    if totalAllowed <= 0:
+                        totalAllowed = 999999
 
                     percentage = float(100) / float(totalAllowed)
-
                     percentage = float(percentage) * float(inMB)
-                except:
-                    return [0,0]
 
-                if percentage > 100.0:
+                    if percentage > 100.0:
-                    percentage = 100
+                        percentage = 100
 
-                return [inMB,percentage]
+                    return [inMB, percentage]
+                except (ValueError, IndexError, IOError):
+                    return [0, 0]
             else:
                 return [0, 0]
 
-
         except OSError as msg:
             logging.CyberCPLogFileWriter.writeToFile(str(msg) + " [findDomainBW]")
-            return 0
+            return [0, 0]
         except ValueError as msg:
             logging.CyberCPLogFileWriter.writeToFile(str(msg) + " [findDomainBW]")
-            return 0
+            return [0, 0]
-
-        return 1
 
     @staticmethod
     def changeSystemLanguage():
+        """Change system language with improved error handling"""
         try:
-
             command = 'localectl set-locale LANG=en_US.UTF-8'
-
             cmd = shlex.split(command)
-
             res = subprocess.call(cmd)
 
             if res == 1:
@@ -135,12 +255,10 @@ class findBWUsage:
             else:
                 pass
 
-
             print("###############################################")
             print("        Language Changed to English                ")
             print("###############################################")
 
-
         except OSError as msg:
             logging.CyberCPLogFileWriter.writeToFile(str(msg) + " [changeSystemLanguage]")
             return 0
@@ -151,4 +269,5 @@ class findBWUsage:
         return 1
 
 
-findBWUsage.startCalculations()
+if __name__ == "__main__":
+    findBWUsage.startCalculations()