import psutil
import time
import csv
import argparse
import os
from datetime import datetime
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from collections import deque
import threading
import sys

class SystemResourceMonitor:
    def __init__(self, interval=1, history_size=100):
        self.interval = interval
        self.history_size = history_size
        self.monitoring = False
        self.data_history = deque(maxlen=history_size)
        
        # 初始化歷史數(shù)據(jù)存儲
        self.timestamps = deque(maxlen=history_size)
        self.cpu_percentages = deque(maxlen=history_size)
        self.memory_percentages = deque(maxlen=history_size)
        self.disk_usages = deque(maxlen=history_size)
        self.network_bytes_sent = deque(maxlen=history_size)
        self.network_bytes_recv = deque(maxlen=history_size)
        
        # 網(wǎng)絡流量基準值
        self.net_io_baseline = psutil.net_io_counters()
        
    def get_cpu_info(self):
        """獲取CPU信息"""
        try:
            cpu_percent = psutil.cpu_percent(interval=1)
            cpu_count = psutil.cpu_count()
            cpu_freq = psutil.cpu_freq()
            
            return {
                'percent': cpu_percent,
                'count': cpu_count,
                'frequency': cpu_freq.current if cpu_freq else 0
            }
        except Exception as e:
            return {'error': str(e)}
            
    def get_memory_info(self):
        """獲取內(nèi)存信息"""
        try:
            memory = psutil.virtual_memory()
            swap = psutil.swap_memory()
            
            return {
                'total': memory.total,
                'available': memory.available,
                'used': memory.used,
                'percent': memory.percent,
                'swap_total': swap.total,
                'swap_used': swap.used,
                'swap_percent': swap.percent
            }
        except Exception as e:
            return {'error': str(e)}
            
    def get_disk_info(self):
        """獲取磁盤信息"""
        try:
            disk = psutil.disk_usage('/')
            disk_io = psutil.disk_io_counters()
            
            return {
                'total': disk.total,
                'used': disk.used,
                'free': disk.free,
                'percent': disk.percent,
                'read_bytes': disk_io.read_bytes if disk_io else 0,
                'write_bytes': disk_io.write_bytes if disk_io else 0
            }
        except Exception as e:
            return {'error': str(e)}
            
    def get_network_info(self):
        """獲取網(wǎng)絡信息"""
        try:
            net_io = psutil.net_io_counters()
            
            # 計算相對于基準值的流量
            bytes_sent = net_io.bytes_sent - self.net_io_baseline.bytes_sent
            bytes_recv = net_io.bytes_recv - self.net_io_baseline.bytes_recv
            
            return {
                'bytes_sent': bytes_sent,
                'bytes_recv': bytes_recv,
                'packets_sent': net_io.packets_sent,
                'packets_recv': net_io.packets_recv
            }
        except Exception as e:
            return {'error': str(e)}
            
    def format_bytes(self, bytes_value):
        """格式化字節(jié)單位"""
        for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
            if bytes_value < 1024.0:
                return f"{bytes_value:.2f} {unit}"
            bytes_value /= 1024.0
        return f"{bytes_value:.2f} PB"
        
    def display_system_info(self):
        """顯示系統(tǒng)信息"""
        # 清屏（兼容不同操作系統(tǒng)）
        os.system('cls' if os.name == 'nt' else 'clear')
        
        print("=" * 80)
        print("系統(tǒng)資源監(jiān)控工具")
        print("=" * 80)
        print(f"監(jiān)控時間: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print()
        
        # CPU信息
        cpu_info = self.get_cpu_info()
        if 'error' not in cpu_info:
            print("CPU信息:")
            print(f"  使用率: {cpu_info['percent']:.1f}%")
            print(f"  核心數(shù): {cpu_info['count']}")
            print(f"  頻率: {cpu_info['frequency']:.2f} MHz")
        else:
            print(f"CPU信息獲取失敗: {cpu_info['error']}")
        print()
        
        # 內(nèi)存信息
        memory_info = self.get_memory_info()
        if 'error' not in memory_info:
            print("內(nèi)存信息:")
            print(f"  總計: {self.format_bytes(memory_info['total'])}")
            print(f"  已用: {self.format_bytes(memory_info['used'])}")
            print(f"  可用: {self.format_bytes(memory_info['available'])}")
            print(f"  使用率: {memory_info['percent']:.1f}%")
            if memory_info['swap_total'] > 0:
                print(f"  交換分區(qū): {memory_info['swap_percent']:.1f}%")
        else:
            print(f"內(nèi)存信息獲取失敗: {memory_info['error']}")
        print()
        
        # 磁盤信息
        disk_info = self.get_disk_info()
        if 'error' not in disk_info:
            print("磁盤信息:")
            print(f"  總計: {self.format_bytes(disk_info['total'])}")
            print(f"  已用: {self.format_bytes(disk_info['used'])}")
            print(f"  可用: {self.format_bytes(disk_info['free'])}")
            print(f"  使用率: {disk_info['percent']:.1f}%")
        else:
            print(f"磁盤信息獲取失敗: {disk_info['error']}")
        print()
        
        # 網(wǎng)絡信息
        network_info = self.get_network_info()
        if 'error' not in network_info:
            print("網(wǎng)絡信息:")
            print(f"  發(fā)送: {self.format_bytes(network_info['bytes_sent'])}")
            print(f"  接收: {self.format_bytes(network_info['bytes_recv'])}")
            print(f"  發(fā)送包: {network_info['packets_sent']:,}")
            print(f"  接收包: {network_info['packets_recv']:,}")
        else:
            print(f"網(wǎng)絡信息獲取失敗: {network_info['error']}")
        print()
        
        print("=" * 80)
        print("按 Ctrl+C 停止監(jiān)控")
        
    def collect_data(self):
        """收集監(jiān)控數(shù)據(jù)"""
        timestamp = datetime.now()
        
        data = {
            'timestamp': timestamp,
            'cpu': self.get_cpu_info(),
            'memory': self.get_memory_info(),
            'disk': self.get_disk_info(),
            'network': self.get_network_info()
        }
        
        # 添加到歷史記錄
        self.data_history.append(data)
        
        # 更新圖表數(shù)據(jù)
        self.timestamps.append(timestamp)
        if 'error' not in data['cpu']:
            self.cpu_percentages.append(data['cpu']['percent'])
        if 'error' not in data['memory']:
            self.memory_percentages.append(data['memory']['percent'])
        if 'error' not in data['disk']:
            self.disk_usages.append(data['disk']['percent'])
            
        return data
        
    def save_to_csv(self, filename="system_monitor.csv"):
        """保存數(shù)據(jù)到CSV文件"""
        try:
            file_exists = os.path.exists(filename)
            
            with open(filename, 'a', newline='', encoding='utf-8') as csvfile:
                fieldnames = [
                    'timestamp', 'cpu_percent', 'memory_percent', 'disk_percent',
                    'network_sent', 'network_recv'
                ]
                writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
                
                # 如果是新文件，寫入表頭
                if not file_exists:
                    writer.writeheader()
                    
                # 寫入最新數(shù)據(jù)
                if self.data_history:
                    latest_data = self.data_history[-1]
                    row = {
                        'timestamp': latest_data['timestamp'].strftime('%Y-%m-%d %H:%M:%S'),
                        'cpu_percent': latest_data['cpu'].get('percent', 0) if 'error' not in latest_data['cpu'] else 0,
                        'memory_percent': latest_data['memory'].get('percent', 0) if 'error' not in latest_data['memory'] else 0,
                        'disk_percent': latest_data['disk'].get('percent', 0) if 'error' not in latest_data['disk'] else 0,
                        'network_sent': latest_data['network'].get('bytes_sent', 0) if 'error' not in latest_data['network'] else 0,
                        'network_recv': latest_data['network'].get('bytes_recv', 0) if 'error' not in latest_data['network'] else 0
                    }
                    writer.writerow(row)
                    
            return True
        except Exception as e:
            print(f"保存CSV文件時出錯: {e}")
            return False
            
    def generate_chart(self, filename="resource_usage.png"):
        """生成資源使用趨勢圖"""
        try:
            if len(self.timestamps) < 2:
                print("數(shù)據(jù)不足，無法生成圖表")
                return False
                
            plt.figure(figsize=(12, 8))
            
            # 創(chuàng)建子圖
            ax1 = plt.subplot(2, 2, 1)
            if self.cpu_percentages:
                ax1.plot(self.timestamps, self.cpu_percentages, 'b-', label='CPU使用率')
                ax1.set_title('CPU使用率')
                ax1.set_ylabel('百分比 (%)')
                ax1.legend()
                ax1.grid(True)
            
            ax2 = plt.subplot(2, 2, 2)
            if self.memory_percentages:
                ax2.plot(self.timestamps, self.memory_percentages, 'g-', label='內(nèi)存使用率')
                ax2.set_title('內(nèi)存使用率')
                ax2.set_ylabel('百分比 (%)')
                ax2.legend()
                ax2.grid(True)
            
            ax3 = plt.subplot(2, 2, 3)
            if self.disk_usages:
                ax3.plot(self.timestamps, self.disk_usages, 'r-', label='磁盤使用率')
                ax3.set_title('磁盤使用率')
                ax3.set_ylabel('百分比 (%)')
                ax3.set_xlabel('時間')
                ax3.legend()
                ax3.grid(True)
            
            ax4 = plt.subplot(2, 2, 4)
            if self.network_bytes_sent and self.network_bytes_recv:
                ax4.plot(self.timestamps, [b/1024/1024 for b in self.network_bytes_sent], 'c-', label='發(fā)送(MB)')
                ax4.plot(self.timestamps, [b/1024/1024 for b in self.network_bytes_recv], 'm-', label='接收(MB)')
                ax4.set_title('網(wǎng)絡流量')
                ax4.set_ylabel('流量 (MB)')
                ax4.set_xlabel('時間')
                ax4.legend()
                ax4.grid(True)
            
            # 格式化時間軸
            for ax in [ax1, ax2, ax3, ax4]:
                ax.xaxis.set_major_formatter(mdates.DateFormatter('%H:%M:%S'))
                ax.xaxis.set_major_locator(mdates.MinuteLocator(interval=5))
                plt.setp(ax.xaxis.get_majorticklabels(), rotation=45)
            
            plt.tight_layout()
            plt.savefig(filename, dpi=300, bbox_inches='tight')
            plt.close()
            
            print(f"圖表已保存到: {filename}")
            return True
        except Exception as e:
            print(f"生成圖表時出錯: {e}")
            return False
            
    def check_thresholds(self, thresholds):
        """檢查閾值并發(fā)出告警"""
        if not self.data_history:
            return
            
        latest_data = self.data_history[-1]
        
        # 檢查CPU閾值
        if 'error' not in latest_data['cpu'] and thresholds.get('cpu'):
            cpu_percent = latest_data['cpu']['percent']
            if cpu_percent > thresholds['cpu']:
                print(f"??  CPU使用率過高: {cpu_percent:.1f}% (閾值: {thresholds['cpu']}%)")
                
        # 檢查內(nèi)存閾值
        if 'error' not in latest_data['memory'] and thresholds.get('memory'):
            memory_percent = latest_data['memory']['percent']
            if memory_percent > thresholds['memory']:
                print(f"??  內(nèi)存使用率過高: {memory_percent:.1f}% (閾值: {thresholds['memory']}%)")
                
        # 檢查磁盤閾值
        if 'error' not in latest_data['disk'] and thresholds.get('disk'):
            disk_percent = latest_data['disk']['percent']
            if disk_percent > thresholds['disk']:
                print(f"??  磁盤使用率過高: {disk_percent:.1f}% (閾值: {thresholds['disk']}%)")
                
    def start_monitoring(self, duration=None, csv_file=None, chart_file=None, thresholds=None):
        """開始監(jiān)控"""
        self.monitoring = True
        start_time = time.time()
        
        try:
            while self.monitoring:
                # 收集數(shù)據(jù)
                self.collect_data()
                
                # 顯示信息
                self.display_system_info()
                
                # 保存到CSV
                if csv_file:
                    self.save_to_csv(csv_file)
                    
                # 檢查閾值
                if thresholds:
                    self.check_thresholds(thresholds)
                    
                # 檢查是否達到指定時長
                if duration and (time.time() - start_time) >= duration:
                    break
                    
                # 等待下次監(jiān)控
                time.sleep(self.interval)
                
        except KeyboardInterrupt:
            print("\n\n停止監(jiān)控...")
        finally:
            self.monitoring = False
            
            # 生成圖表
            if chart_file:
                self.generate_chart(chart_file)
                
            print("監(jiān)控已停止")

def main():
    parser = argparse.ArgumentParser(description="系統(tǒng)資源監(jiān)控工具")
    parser.add_argument("-i", "--interval", type=int, default=5, 
                       help="監(jiān)控間隔（秒，默認：5）")
    parser.add_argument("-d", "--duration", type=int, 
                       help="監(jiān)控時長（秒，不指定則持續(xù)監(jiān)控）")
    parser.add_argument("-c", "--csv", help="保存數(shù)據(jù)到CSV文件")
    parser.add_argument("-g", "--graph", help="生成圖表文件")
    parser.add_argument("--cpu-threshold", type=float, 
                       help="CPU使用率告警閾值（%）")
    parser.add_argument("--memory-threshold", type=float, 
                       help="內(nèi)存使用率告警閾值（%）")
    parser.add_argument("--disk-threshold", type=float, 
                       help="磁盤使用率告警閾值（%）")
    
    args = parser.parse_args()
    
    # 檢查依賴
    try:
        import psutil
    except ImportError:
        print("錯誤: 缺少psutil庫，請安裝: pip install psutil")
        sys.exit(1)
        
    if args.graph and not plt:
        print("警告: 缺少matplotlib庫，無法生成圖表")
        args.graph = None
        
    # 設置閾值
    thresholds = {}
    if args.cpu_threshold:
        thresholds['cpu'] = args.cpu_threshold
    if args.memory_threshold:
        thresholds['memory'] = args.memory_threshold
    if args.disk_threshold:
        thresholds['disk'] = args.disk_threshold
        
    try:
        monitor = SystemResourceMonitor(interval=args.interval)
        monitor.start_monitoring(
            duration=args.duration,
            csv_file=args.csv,
            chart_file=args.graph,
            thresholds=thresholds
        )
        
    except Exception as e:
        print(f"程序執(zhí)行出錯: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()

pip install psutil matplotlib

# 基本用法，每5秒刷新一次
python system_monitor.py

# 設置監(jiān)控間隔為10秒
python system_monitor.py -i 10

# 監(jiān)控60秒后自動停止
python system_monitor.py -d 60

# 保存監(jiān)控數(shù)據(jù)到CSV文件
python system_monitor.py -c monitor_data.csv

# 生成資源使用趨勢圖
python system_monitor.py -g resource_trend.png

# 設置告警閾值
python system_monitor.py --cpu-threshold 80 --memory-threshold 85 --disk-threshold 90

python system_monitor.py

python system_monitor.py -i 10 -d 300 -c system_data.csv -g trend.png

python system_monitor.py --cpu-threshold 80 --memory-threshold 85