import sys
import asyncio
import subprocess
from queue import Queue
from PyQt5.QtCore import QThread, pyqtSignal, QTimer
from PyQt5.QtGui import QColor, QFont
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QHBoxLayout, QPushButton, QLabel, QLineEdit, QTextEdit, QProgressBar, QTableWidget, QTableWidgetItem, QMessageBox, QFileDialog, QInputDialog

class PingThread(QThread):
    # 定義信號
    result_signal = pyqtSignal(str, str)  # 傳遞IP和狀態(tài)
    progress_signal = pyqtSignal(int)  # 用于更新進度條的信號
    stop_signal = pyqtSignal()  # 用于停止線程的信號

    def __init__(self, ip_queue, total_ips, parent=None):
        super().__init__(parent)
        self.ip_queue = ip_queue
        self.total_ips = total_ips
        self.pinged_ips = 0  # 已ping的IP數(shù)量
        self.stop_requested = False  # 停止請求標志位
        self.online_ips = []  # 在線IP列表
        self.offline_ips = []  # 掉線IP列表

    def run(self):
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        loop.run_until_complete(self.ping_ips())

    async def ping_ips(self):
        tasks = []
        while not self.ip_queue.empty() and not self.stop_requested:
            ip = self.ip_queue.get()
            task = asyncio.create_task(self.ping(ip))
            tasks.append(task)
            if len(tasks) >= 20:  # 每次最多并發(fā)20個請求
                await asyncio.gather(*tasks)
                tasks.clear()
        if tasks:
            await asyncio.gather(*tasks)

    async def ping(self, ip):
        res = await asyncio.to_thread(self.ping_single, ip)
        status = "在線" if res == 0 else "掉線"
        self.result_signal.emit(ip, status)  # 發(fā)送IP和狀態(tài)信號

        # 根據(jù)狀態(tài)將IP添加到對應的列表
        if status == "在線":
            self.online_ips.append(ip)
        else:
            self.offline_ips.append(ip)

        self.pinged_ips += 1
        progress = int((self.pinged_ips / self.total_ips) * 100)  # 計算當前進度
        self.progress_signal.emit(progress)  # 發(fā)送進度信號

    def ping_single(self, ip):
        """使用subprocess調(diào)用ping命令"""
        return subprocess.call(f'ping -n 1 -w 5 {ip}', stdout=subprocess.PIPE)

    def stop(self):
        self.stop_requested = True  # 設(shè)置停止標志


class PingApp(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("IP網(wǎng)絡(luò)掃描工具")
        self.setGeometry(200, 200, 800, 600)

        # 設(shè)置最小寬度為1064像素
        self.setMinimumWidth(1064)

        # 設(shè)置全局字體為Segoe UI
        font = QFont("Segoe UI", 10)
        self.setFont(font)

        # UI控件
        self.layout = QVBoxLayout()

        self.status_label = QLabel("請輸入待檢測的網(wǎng)段列表 (例如：192.168.1, 192.168.2,支持多網(wǎng)段掃描，網(wǎng)段之間用英文','隔開)", self)
        self.layout.addWidget(self.status_label)

        self.input_field = QLineEdit(self)
        self.layout.addWidget(self.input_field)

        # 創(chuàng)建水平布局來放置開始、停止和導出按鈕
        self.button_layout = QHBoxLayout()

        # 修改按鈕背景色為淡藍色
        self.start_button = QPushButton("開始掃描", self)
        self.start_button.setStyleSheet("background-color: lightblue;")
        self.start_button.clicked.connect(self.start_ping)
        self.button_layout.addWidget(self.start_button)

        self.stop_button = QPushButton("停止掃描", self)
        self.stop_button.setStyleSheet("background-color: lightblue;")
        self.stop_button.clicked.connect(self.stop_ping)
        self.stop_button.setEnabled(False)  # 初始時不可用
        self.button_layout.addWidget(self.stop_button)

        self.export_button = QPushButton("導出數(shù)據(jù)", self)
        self.export_button.setStyleSheet("background-color: lightblue;")
        self.export_button.clicked.connect(self.export_data)
        self.button_layout.addWidget(self.export_button)

        # 將按鈕布局添加到主布局中
        self.layout.addLayout(self.button_layout)

        # 添加文本輸出區(qū)域
        self.output_text = QTextEdit(self)
        self.layout.addWidget(self.output_text)

        # 設(shè)置最大高度來縮減文本顯示區(qū)域的高度
        self.output_text.setMaximumHeight(200)  # 設(shè)置最大高度為200

        # 添加進度條
        self.progress_bar = QProgressBar(self)
        self.progress_bar.setRange(0, 100)
        self.layout.addWidget(self.progress_bar)

        # 添加表格
        self.table_widget = QTableWidget(self)
        self.layout.addWidget(self.table_widget)

        self.setLayout(self.layout)

    def start_ping(self):
        user_input = self.input_field.text().strip()
        if not user_input:
            self.output_text.append("輸入無效，請輸入網(wǎng)段列表！")
            return

        # 校驗輸入是否是有效的網(wǎng)段格式
        if not self.is_valid_ip_prefix(user_input):
            self.output_text.append("無效的網(wǎng)段格式，請確保輸入的是有效網(wǎng)段")
            return

        self.output_text.append(f"開始檢測：{user_input}")
        self.start_button.setEnabled(False)
        self.stop_button.setEnabled(True)  # 啟用停止按鈕

        ip_queue = Queue()
        ip_prefixes = user_input.split(',')  # 支持多個網(wǎng)段
        total_ips = len(ip_prefixes) * 256  # 每個網(wǎng)段包含256個IP地址

        # 遍歷多個網(wǎng)段并生成IP地址
        for prefix in ip_prefixes:
            prefix = prefix.strip()
            for i in range(256):
                ip = f"{prefix}.{i}"  # 生成每個網(wǎng)段對應的IP
                ip_queue.put(ip)

        # 動態(tài)設(shè)置表格的行數(shù)和列數(shù)
        rows = (total_ips + 9) // 10  # 計算需要的行數(shù)，最多每行10個IP
        self.table_widget.setRowCount(rows)  # 設(shè)置表格的行數(shù)
        self.table_widget.setColumnCount(10)  # 每行10個IP

        # 填充表格
        ip_index = 0
        for prefix in ip_prefixes:  # 遍歷每個網(wǎng)段
            prefix = prefix.strip()
            for i in range(256):
                if ip_index < total_ips:
                    ip = f"{prefix}.{i}"  # 生成完整的IP地址
                    row = ip_index // 10  # 計算行數(shù)
                    col = ip_index % 10  # 計算列數(shù)
                    self.table_widget.setItem(row, col, QTableWidgetItem(ip))  # 設(shè)置表格單元格的IP
                    ip_index += 1

        # 啟動ping線程
        self.ping_thread = PingThread(ip_queue, total_ips)
        self.ping_thread.result_signal.connect(self.update_result)
        self.ping_thread.progress_signal.connect(self.update_progress)  # 連接進度更新信號
        self.ping_thread.finished.connect(self.on_ping_finished)
        self.ping_thread.start()

    def stop_ping(self):
        if self.ping_thread.isRunning():
            self.ping_thread.stop()  # 請求停止線程
            self.output_text.append("檢測已停止")
            self.start_button.setEnabled(True)
            self.stop_button.setEnabled(False)  # 停止后禁用停止按鈕

    def export_data(self):
        # 使用 QInputDialog.getItem() 來顯示選擇框
        export_type, ok = QInputDialog.getItem(self, "選擇導出項", "選擇導出數(shù)據(jù)項", 
                                               ["在線IP", "掉線IP", "日志記錄"], 0, False)

        if ok:
            if export_type == "在線IP":
                self.export_file("在線IP.txt")
            elif export_type == "掉線IP":
                self.export_file("掉線IP.txt")
            elif export_type == "日志記錄":
                self.export_log()

    def export_file(self, file_name):
        # 彈出文件保存對話框，選擇保存路徑
        file_path, _ = QFileDialog.getSaveFileName(self, f"保存{file_name}", file_name, "文本文件 (*.txt)")
        if file_path:
            # 使用內(nèi)存列表替代讀取文件
            ip_list = self.get_ip_list(file_name)

            # 排序IP地址（從小到大）
            ip_list.sort(key=self.ip_to_tuple)

            # 寫入排序后的文件，每個IP地址按行分開
            with open(file_path, 'w') as f:
                for ip in ip_list:
                    f.write(f"{ip}\n")  # 每個IP地址后加上換行符
            QMessageBox.information(self, "導出成功", f"{file_name} 已成功導出到 {file_path}")

    def export_log(self):
        # 彈出文件保存對話框，選擇保存日志路徑
        file_path, _ = QFileDialog.getSaveFileName(self, "保存日志文件", "日志記錄.txt", "文本文件 (*.txt)")
        if file_path:
            with open(file_path, 'w') as f:
                f.write(self.output_text.toPlainText())
            QMessageBox.information(self, "導出成功", f"日志記錄已成功導出到 {file_path}")

    def get_ip_list(self, file_name):
        # 根據(jù)文件名從內(nèi)存中獲取相應的IP列表
        ip_list = []
        if file_name == "在線IP.txt":
            ip_list = self.ping_thread.online_ips  # 獲取在線IP列表
        elif file_name == "掉線IP.txt":
            ip_list = self.ping_thread.offline_ips  # 獲取掉線IP列表
        return ip_list

    def ip_to_tuple(self, ip):
        # 將IP地址轉(zhuǎn)換為元組，用于排序
        return tuple(map(int, ip.split('.')))

    def is_valid_ip_prefix(self, ip_prefix):
        # 簡單的IP前綴校驗（支持多個網(wǎng)段）
        parts = ip_prefix.split(',')
        if len(parts) < 1:
            return False
        for part in parts:
            part = part.strip()
            if not self.is_valid_single_ip_prefix(part):
                return False
        return True

    def is_valid_single_ip_prefix(self, ip_prefix):
        # 校驗單一網(wǎng)段是否有效（例如：192.168.1）
        parts = ip_prefix.split('.')
        if len(parts) != 3:
            return False
        for part in parts:
            if not part.isdigit() or not 0 <= int(part) <= 255:
                return False
        return True  # 必須輸入三個部分

    def update_result(self, ip, status):
        # 更新表格單元格的背景色
        for row in range(self.table_widget.rowCount()):
            for col in range(self.table_widget.columnCount()):
                item = self.table_widget.item(row, col)
                if item and item.text() == ip:
                    if status == "在線":
                        item.setBackground(QColor('green'))  # 使用QColor類
                    else:
                        item.setBackground(QColor('gray'))  # 使用QColor類

        # 在文本輸出區(qū)域顯示每個IP的狀態(tài)
        self.output_text.append(f"{ip} - {status}")

    def update_progress(self, progress):
        self.progress_bar.setValue(progress)  # 更新進度條的值

    def on_ping_finished(self):
        # 計算在線和掉線主機數(shù)量
        online_count = len(self.ping_thread.online_ips)
        offline_count = len(self.ping_thread.offline_ips)

        # 顯示檢測完成信息和統(tǒng)計
        self.output_text.append(f"掃描完成！")
        self.output_text.append(f"在線主機數(shù): {online_count}")
        self.output_text.append(f"掉線主機數(shù): {offline_count}")

        self.start_button.setEnabled(True)
        self.stop_button.setEnabled(False)  # 禁用停止按鈕
        # 彈出完成提示框
        QMessageBox.information(self, "提示", "掃描已完成！")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = PingApp()
    window.show()
    sys.exit(app.exec_())