NetWork_v1/SimObjs.py

import math
import sys
import threading
from time import sleep
from tkinter import Tk
from tkinter import messagebox

from ttkbootstrap import *
from uuid import uuid4
import ipaddress
import time

from PIL import ImageTk, Image, ImageFont

from dbUtil import search
from packet import *

pause_event = threading.Event()

class SimBase():
    # todo: 组件父类
    """
    图标类，所有组件的父类
    """
    def __init__(self, canvas, x, y, id, config=None, label=None):
        self.ConfigCorrect = 0 if config is None else config  # 是否进行配置
        self.ObjType = None  # 组件类型 1->主机，2->路由器，3->交换机，4->集线器
        self.ObjLabel = ""
        self.ObjID = id
        self.ObjX = x
        self.ObjY = y
        self.canvas = canvas
        self.is_packet = False
        self.is_un_packet = False
        self.host_img = ImageTk.PhotoImage(
            Image.open(sys.path[0] + "/../datas/images/主机@2x.png").resize((40, 40)))
        self.host_img_tm = ImageTk.PhotoImage(
            Image.open(sys.path[0] + "/../datas/images/主机--暗色@2x.png").resize((40, 40)))
        self.router_img = ImageTk.PhotoImage(
            Image.open(sys.path[0] + "/../datas/images/路由器@2x.png").resize((40, 40)))
        self.router_img_tm = ImageTk.PhotoImage(
            Image.open(sys.path[0] + "/../datas/images/路由器--暗色@2x.png").resize((40, 40)))
        self.switch_img = ImageTk.PhotoImage(
            Image.open(sys.path[0] + "/../datas/images/交换机@2x.png").resize((40, 40)))
        self.switch_img_tm = ImageTk.PhotoImage(
            Image.open(sys.path[0] + "/../datas/images/交换机--暗色@2x.png").resize((40, 40)))
        self.hub_img = ImageTk.PhotoImage(
            Image.open(sys.path[0] + "/../datas/images/集线器@2x.png").resize((40, 40)))
        self.hub_img_tm = ImageTk.PhotoImage(
            Image.open(sys.path[0] + "/../datas/images/集线器--暗色@2x.png").resize((40, 40)))
        self.img = None
        self.img_tm = None
        self.interface = [{}, {}, {}, {}]
        self.connections = ["1", "2", "3", "4"]
        self.set_default_config()
        self.packet_window = None

    def set_default_name(self):
        data_frame = search(f"select objid from sim_objs where objType={self.ObjType}")
        num = data_frame.size + 1
        if isinstance(self, SimHost):
            name = "SHO%d" % num
        elif isinstance(self, SimRouter):
            name = "SRO%d" % num
        elif isinstance(self, SimSwitch):
            name = "SWI%d" % num
        else:
            name = "SHUB%d" % num
        return name

    def create_img(self):
        """
        创建图片
        :return:
        """
        if self.ObjType == 1:
            self.img = self.host_img
            self.img_tm = self.host_img_tm
        elif self.ObjType == 2:
            self.img = self.router_img
            self.img_tm = self.router_img_tm
        elif self.ObjType == 3:
            self.img = self.switch_img
            self.img_tm = self.switch_img_tm
        else:
            self.img = self.hub_img
            self.img_tm = self.hub_img_tm
        self.canvas.delete("L")
        id = self.canvas.create_image(self.ObjX - 30, self.ObjY - 30,
                                      image=self.img if self.ConfigCorrect == 1 else self.img_tm, anchor="nw",
                                      tags=self.ObjID)
        self.canvas.dtag("L", "L")
        self.canvas.create_text(self.ObjX - 20, self.ObjY - 60, text=self.ObjLabel, font=("", 16, "bold"),
                                fill="#9b78eb", tags=self.ObjID + "text", anchor="nw")
        if self.ObjType == 1:
            self.canvas.create_text(self.ObjX - 45, self.ObjY + 15, text=self.interface[0]["ip"]if len(self.interface) > 0 else "", font=("", 16, "bold"),
                                    fill="#9b78eb", tags=self.ObjID + "text", anchor="nw")
            self.canvas.create_text(self.ObjX - 45, self.ObjY + 35, text=self.interface[0]["mac"]if len(self.interface) > 0 else "", font=("", 16, "bold"),
                                    fill="#9b78eb", tags=self.ObjID + "text", anchor="nw")
            self.canvas.create_text(self.ObjX - 45, self.ObjY + 55, text=self.interface[0]["addr"]if len(self.interface) > 0 else "", font=("", 16, "bold"),
                                    fill="#9b78eb", tags=self.ObjID + "text", anchor="nw")
        self.canvas.tag_raise(id)

    def set_default_config(self):
        sql = f"select * from conn_config where node_id='{self.ObjID}'"
        conn_data = search(sql)
        for index, conn in conn_data.iterrows():
            self.interface[int(conn["node_ifs"]) - 1] = {"ip": conn["ip"],
                                                         "mac": conn["mac"],
                                                         "conn_port": conn["conn_port"],
                                                         "addr": conn["addr"]}

    def get_config(self):
        sql = f"select * from conn_config where node_id='{self.ObjID}'"
        conn_data = search(sql)
        return conn_data

    def transfer_animate(self, status, packet, packet_size, error_message=None):
        if status:
            text = f"消息大小: {packet_size}\n" \
                   f"消息内容: {packet}"
            self.canvas.create_rectangle(self.ObjX + 30, self.ObjY - 30, self.ObjX + 160, self.ObjY + 20,
                                         outline="#92d050",
                                         width=3, fill="#92d050", tags=self.ObjID + "packetData")
            self.canvas.create_text(self.ObjX + 35, self.ObjY - 25, text=text, anchor="nw",
                                    font=('', 10), tags=self.ObjID + "packetData")  # 显示文字
            self.canvas.update()
            sleep(2)
            self.canvas.delete(self.ObjID + "packetData")  # 删除展示的数据包内容
        else:
            text = f"传输失败\n" if error_message is None else error_message
            self.canvas.create_rectangle(self.ObjX + 30, self.ObjY - 30, self.ObjX + 160, self.ObjY, outline="red",
                                         width=3, fill="red", tags=self.ObjID + "packetData")
            self.canvas.create_text(self.ObjX + 35, self.ObjY - 25, text=text, anchor="nw",
                                    font=('', 10), tags=self.ObjID + "packetData")  # 显示文字
            self.canvas.update()
            sleep(2)
            self.canvas.delete(self.ObjID + "packetData")  # 删除展示的数据包内容

    def create_detail_button(self):
        frist_x = self.ObjX  # 获取frist_x
        frist_y = self.ObjY  # 获取frist_y
        tag_id = self.ObjID + "detail_button"
        self.canvas.create_rectangle(frist_x + 10, frist_y - 25, frist_x + 40, frist_y - 15, fill="#f0a732",
                                     outline="#47b2ff",width=2, tags=tag_id)

    def __str__(self):
        str = ""
        config = self.get_config()
        for index, data in config.iterrows():
            str += f"【接口{data['node_ifs']}】\n"
            str += f"IP: {data['ip']}\n"
            str += f"MAC: {data['mac']}\n"
        return str


class SimPacket():
    # todo: 数据包类
    """
    数据包类
    """
    def __init__(self, source_ip, source_mac, destination_ip, destination_mac, message, label, size, sendObj, receiveObj):
        """
        :param source_mac: 源主机mac地址
        :param source_ip: 源主机ip地址
        :param destination_mac: 目的主机mac地址
        :param destination_ip: 目的主机ip地址
        :param message: 数据
        """
        self.source_mac = source_mac
        self.source_ip = source_ip
        self.sendObj = sendObj
        self.receiveObj = receiveObj
        self.destination_mac = destination_mac
        self.destination_ip = destination_ip
        self.message = message
        self.label = label
        self.size = size
        self.up_jump = None  # 上一跳连接对象
        self.objs = set()
        self.img = None
        self.id = str(uuid4())
        self.process_img()

    def process_img(self):
        img = Image.open(sys.path[0] + "/../datas/images/packet.png").resize((30, 30))
        draw = ImageDraw.Draw(img)
        font = ImageFont.truetype("arial.ttf", size=16)
        text_width, text_height = draw.textsize(str(self.label), font=font)
        # 计算文本在右上角的位置
        x = img.width - text_width - 10  # 右上角，距离右边缘10像素
        y = 10  # 距离上边缘10像素
        # 绘制文本
        draw.text((x, y), str(self.label), font=font, fill=(255, 0, 0))  # 红色字体
        self.img = ImageTk.PhotoImage(img)

    def move(self, id,  cv, target_x, target_y, duration):
        start_x, start_y = cv.coords(id)
        distance_x = target_x - start_x
        distance_y = target_y - start_y
        steps = duration // 10  # 以10毫秒为间隔进行移动
        step_x = distance_x / steps
        step_y = distance_y / steps
        self._move_step(id, cv, start_x, start_y, target_x, target_y, step_x, step_y, steps)

    def _move_step(self,id, cv, start_x, start_y, target_x, target_y, step_x, step_y, steps):
        if pause_event.is_set():
            new_x = start_x + step_x
            new_y = start_y + step_y
            sleep(2)
            self._move_step(id, cv, new_x, new_y, target_x, target_y, step_x, step_y, steps)
        if steps > 0:
            new_x = start_x + step_x
            new_y = start_y + step_y
            cv.coords(id, new_x, new_y)
            cv.update()  # 更新画布显示
            sleep(0.01)  # 添加延迟以控制动画速度
            self._move_step(id, cv, new_x, new_y, target_x, target_y, step_x, step_y, steps - 1)
        else:
            cv.coords(id, target_x, target_y)
            cv.delete(id)

    def transfer_packet(self, cv: Canvas, nodex_tag: SimBase, nodey_tag: SimBase):
        self.objs.add(nodex_tag)
        self.objs.add(nodey_tag)
        cv.create_image(nodex_tag.ObjX - 15, nodex_tag.ObjY - 15, image=self.img, anchor="nw", tags=self.id)
        self.move(self.id, cv, nodey_tag.ObjX - 15, nodey_tag.ObjY - 15, 500)


class AllSimConnect():
    # todo: 连接类
    def __init__(self, canvas, nodex: SimBase, nodex_ifs, nodey: SimBase, nodey_ifs, config=None):
        """
        连接对象
        :param nodex: 节点
        :param nodex_ifs: 节点接口
        :param nodey: 节点
        :param nodey_ifs: 节点接口
        """
        self.canvas = canvas
        self.ConfigCorrect = 0 if config is None else config
        self.NobjS = nodex
        self.NobjE = nodey
        self.IfsS = nodex_ifs
        self.IfsE = nodey_ifs
        self.width = False if self.ConfigCorrect == 0 else True  # 线的粗细

    def transfer(self, source_node, packet: SimPacket):
        """
        传输数据
        :param packet: 数据包
        :return:
        """
        if source_node == self.NobjS:
            packet.transfer_packet(self.canvas, self.NobjS, self.NobjE)
            self.NobjE.receive(packet)
        else:
            packet.transfer_packet(self.canvas, self.NobjE, self.NobjS)
            self.NobjS.receive(packet)


    def draw_line(self):
        if self.width:
            line = self.canvas.create_line(self.NobjS.ObjX, self.NobjS.ObjY, self.NobjE.ObjX, self.NobjE.ObjY,
                                           width=2, fill="#c372f0", tags=self.NobjS.ObjID + self.NobjE.ObjID + "line")
        else:
            line = self.canvas.create_line(self.NobjS.ObjX, self.NobjS.ObjY, self.NobjE.ObjX, self.NobjE.ObjY,
                                           width=2, dash=(250, 4), fill="#c372f0",
                                           tags=self.NobjS.ObjID + self.NobjE.ObjID + "line")
        self.canvas.tag_lower(line, 2)


class SimHost(SimBase):
    # todo: 主机类
    """
    主机类
    """

    def __init__(self, canvas, x, y, id=None, config=None, label=None):
        self.ObjID = str(uuid4()) if id is None else id
        super().__init__(canvas, x, y, self.ObjID, config, label)
        self.ObjType = 1
        self.ObjLabel = label if label is not None else self.set_default_name()
        self.interface = [{}]
        self.connections = [None]
        self.set_default_config()

    def packet_ok(self, receiveObj, label, size):
        count = split_appdata(int(size))
        for i in range(2, count + 1):
            pack = SimPacket(self.interface[0]["ip"], self.interface[0]["mac"],
                             receiveObj.interface[0]["ip"], receiveObj.interface[0]["mac"], label, i, size, self,
                             receiveObj)
            threading.Timer(1 * (i - 1), function=self.send, args=(pack,), ).start()

    def create_packet(self, receiveObj, label, size):
        """
        创建数据包
        :param receiveObj: 目的主机对象
        :param message: 消息
        :return:
        """
        def send(message):
            print(message)
            pack = SimPacket(self.interface[0]["ip"], self.interface[0]["mac"],
                             receiveObj.interface[0]["ip"], receiveObj.interface[0]["mac"], message, 1, size, self,
                             receiveObj)
            self.send(pack)
        message = """发送主机：{}\n\nIP地址：{}\n\nMac地址：{}\n\n发送数据包大小：{}\n\n已发送数据数量：{}\n\n需要发送的数据包总数：{}"""
        self.packet_window = PacketWindow(self.canvas, packet={"packet": True,
                                                                      "size": int(size),
                                                                      "sendObj": self,
                                                                      "receiveObj": receiveObj,
                                                                      "tag": label,
                                                                      "send": send,
                                                                      "mode": [{"app": True},{"trans": True},{"ip": True},{"mac": True}]})
        self.packet_window.title("{} 封包".format(self.ObjLabel))

    def send(self, packet):
        """
        发送数据包
        :param packet:
        :return:
        """
        connection: AllSimConnect = self.connections[0]
        print(f"数据包从 {self.ObjLabel} 发出")
        packet.up_jump = connection
        connection.transfer(self, packet)

    def receive(self, packet: SimPacket):
        """
        接收数据
        :param packet: 数据包
        :return:
        """
        data = ""
        def receive(message):
            data = message
            self.is_un_packet = True
            packet.sendObj.packet_ok(self, data, int(packet.size))
        if not self.is_un_packet:
            PacketWindow(self.canvas, {"packet": False,
                                      "size": None,
                                      "sendObj": packet.sendObj,
                                      "receiveObj": self,
                                      "tag": packet.message,
                                      "send": receive,
                                      "mode": [{"app": True},{"trans": True},{"ip": True},{"mac": True}]}, packet_step=["app", "trans", "ip", "mac"][::-1])
        size = int(packet.size)
        message = """接收主机：{}\n\nIP地址：{}\n\nMac地址：{}\n\n发送数据包大小：{}\n\n已接收数据数量：{}\n\n需要发送的数据包总数：{}"""
        self.canvas.show_bottom_message(self, message=message.format(self.ObjLabel, self.interface[0]["ip"],
                                                                     self.interface[0]["mac"],
                                                                     size, packet.label, split_appdata(size)))
        if split_appdata(size) == packet.label:
            if packet.destination_ip == self.interface[0]["ip"]:
                self.transfer_animate(True, packet.message, size)
                self.canvas.message.show_message(f"{self.ObjLabel}接收到数据：{data}")
            else:
                self.transfer_animate(False, data, size)
            self.canvas.trans_over(packet.objs)


    def __str__(self):
        str = ""
        config = self.get_config()
        for index, data in config.iterrows():
            str += f"【接口{data['node_ifs']}】\n"
            str += f"AppAddr: /Root\n"
            str += f"PORT: {data['conn_port']}\n"
            str += f"IP: {data['ip']}\n"
            str += f"MAC: {data['mac']}\n"
        return str


class SimRouter(SimBase):
    # todo: 路由类
    """
    路由类
    """
    def __init__(self, canvas, x, y, id=None, config=None, label=None, *args):
        self.ObjID = str(uuid4()) if id is None else id
        super().__init__(canvas, x, y, self.ObjID, config, label)
        self.ObjType = 2
        self.ObjLabel = label if label is not None else self.set_default_name()
        self.router_table = {}
        self.set_default_router_table()

    def set_default_router_table(self):
        """
        将数据库中的路由表信息提取
        :return:
        """
        sql = f"select * from router_table where obj_id='{self.ObjID}'"
        router_tables = search(sql)
        for index, router_table in router_tables.iterrows():
            if router_table["node_ifs"] in self.router_table:
                self.router_table[router_table["node_ifs"]].append(router_table["segment"])
            else:
                self.router_table[router_table["node_ifs"]] = [router_table["segment"]]

    def check_destination_ip(self, destination_ip, network):
        """
        检查目标ip是否属于网段范围内
        :param destination_ip: 目标ip
        :param network: 网段
        :return:10.2.3.0/24
        """
        ip = ipaddress.ip_address(destination_ip)
        network = ipaddress.ip_network(network)
        if ip in network:
            return True
        if network == "0.0.0.0/24":  # 如果网段为 0.0.0.0/24 则为默认路由
            return True

    def transmit(self, packet: SimPacket):
        """
        转发数据包
        :return:
        """
        def transfer():
            flag = False
            next_hop_ifs = None
            for conn in self.connections:
                if isinstance(conn, AllSimConnect):
                    if conn.ConfigCorrect == 0:
                        continue
                    if conn == packet.up_jump:
                        continue
                    ifs = self.connections.index(conn) + 1
                    if self.router_table.get(ifs) is None:
                        continue
                    for network in self.router_table.get(ifs):
                        if self.check_destination_ip(packet.destination_ip, network):
                            flag = True
                            next_hop_ifs = ifs
            if flag:
                conn = self.connections[next_hop_ifs - 1]
                packet.up_jump = conn
                conn.transfer(self, packet)
            else:
                for conn in self.connections:
                    if isinstance(conn, AllSimConnect):
                        if conn == packet.up_jump:
                            continue
                        if conn.NobjS != self:
                            if conn.NobjE.ObjType == 1:
                                conn.transfer(self, packet)
                                break
                error_message = "路由寻址失败"
                self.transfer_animate(False, packet, error_message)
        def receive(message):
            packet.message = message
            self.is_packet = True
            transfer()
        if not self.is_packet:
            PacketWindow(self.canvas, {"packet": True,
                                       "size": None,
                                       "sendObj": packet.sendObj,
                                       "receiveObj": packet.receiveObj,
                                       "tag": packet.message,
                                       "send": receive,
                                       "mode": [{"app": False}, {"trans": False}, {"ip": True}, {"mac": True}]},
                         packet_step=["ip", "mac"])
        else:
            transfer()

    def receive(self, packet):
        """
        接收数据
        :param packet: 数据包
        :return:
        """
        def receive(message):
            print(message)
            self.is_un_packet = True
            packet.message = message
            sleep(1.5)
            self.transmit(packet)
        if not self.is_un_packet:
            PacketWindow(self.canvas, {"packet": False,
                                       "size": None,
                                       "sendObj": packet.sendObj,
                                       "receiveObj": packet.receiveObj,
                                       "tag": packet.message,
                                       "send": receive,
                                       "mode": [{"app": False}, {"trans": False}, {"ip": True}, {"mac": True}]},
                         packet_step=["ip", "mac"][::-1])
        else:
            self.transmit(packet)


class SimSwitch(SimBase):
    # todo: 交换机类
    """
    交换机类
    """

    def __init__(self, canvas, x, y, id=None, config=None, label=None, *args):
        self.ObjID = str(uuid4()) if id is None else id
        super().__init__(canvas, x, y, self.ObjID, config, label)
        self.ObjType = 3
        self.ObjLabel = label if label is not None else self.set_default_name()
        self.mac_table = {}
        self.set_default_mac_table()

    def set_default_mac_table(self):
        """
        将数据库中的路由表信息提取
        :return:
        """
        sql = f"select * from mac_table where obj_id='{self.ObjID}'"
        router_tables = search(sql)
        for index, router_table in router_tables.iterrows():
            if router_table["node_ifs"] in self.mac_table:
                self.mac_table[router_table["node_ifs"]].append(router_table["mac"])
            else:
                self.mac_table[router_table["node_ifs"]] = [router_table["mac"]]

    def transmit(self, packet: SimPacket):
        """
        转发数据包
        :return:
        """
        flag = False
        next_hub_ifs = None
        for conn in self.connections:
            if isinstance(conn, AllSimConnect):
                if conn.ConfigCorrect == 0:
                    continue
                ifs = self.connections.index(conn) + 1
                if packet.destination_mac in self.mac_table.get(ifs, []):
                    flag = True
                    next_hub_ifs = ifs
        if flag:
            conn = self.connections[next_hub_ifs - 1]
            packet.up_jump = conn
            conn.transfer(self, packet)
            return
        for conn in self.connections:  # 将数据包往所有接口进行转发
            if isinstance(conn, AllSimConnect):
                if conn == packet.up_jump:
                    continue
                if conn.ConfigCorrect == 0:
                    continue
                new_packet = SimPacket(packet.source_ip,
                                       packet.source_mac,
                                       packet.destination_ip,
                                       packet.destination_mac,
                                       packet.message)
                new_packet.up_jump = conn
                threading.Thread(target=conn.transfer, args=(self, new_packet)).start()

    def receive(self, packet: SimPacket):
        """
        接收数据
        :param packet: 数据包
        :return:
        """
        print(f"交换机{self.ObjLabel}接受到数据{packet.message}")
        self.transmit(packet)


class SimHub(SimBase):
    """
    集线器类
    """

    def __init__(self, canvas, x, y, id=None, config=None, label=None, *args):
        self.ObjID = str(uuid4()) if id is None else id
        super().__init__(canvas, x, y, self.ObjID, config, label)
        self.ObjType = 4
        self.ObjLabel = label if label is not None else self.set_default_name()

    def transmit(self, packet: SimPacket):
        """
        集线器转发数据包
        :return:
        """
        for conn in self.connections:  # 将数据包往所有接口进行转发
            if isinstance(conn, AllSimConnect):
                if conn == packet.up_jump:
                    continue
                if conn.ConfigCorrect == 0:
                    continue
                new_packet = SimPacket(packet.source_ip,
                                       packet.source_mac,
                                       packet.destination_ip,
                                       packet.destination_mac,
                                       packet.message)
                new_packet.up_jump = conn
                threading.Thread(target=conn.transfer, args=(self, new_packet)).start()

    def receive(self, packet: SimPacket):
        """
        接收数据
        :param packet: 数据包
        :return:
        """
        print(f"集线器-{self.ObjLabel}接受到数据，将进行转发！")
        self.transmit(packet)