用Python写一个获取主机信息并发送接受的小工具 简易agent

1.server端

1.项目目录

大家都知道Python简单易学习并且生态完善,大部分功能可以直接调库京今天就给大家带来一个小案例 括弧(本人无聊写的)

这里直接上服务端代码 这里我先来放我的项目树

这里lib包放了点数据类，配置文件类，环境变量类，加密类

2.lib包 文件

dataclass.py

import datetime
from sqlalchemy import create_engine, Column, Integer, String, Float, JSON, desc, Boolean, DateTime, text
from sqlalchemy.orm import declarative_base
from sqlalchemy.orm import sessionmaker

from lib.ENV import SQLITE_PATH

Base = declarative_base()


class Machine(Base):
    __tablename__ = 'machines'
    id = Column(Integer, primary_key=True)
    hostname = Column(String(255))
    cpu_info = Column(JSON)
    interfaces = Column(JSON)
    memory_info = Column(JSON)
    disk_info = Column(JSON)
    disk_uses = Column(JSON)
    processes = Column(JSON)
    users = Column(JSON)
    status = Column(Boolean, default=False)
    datetime = Column(DateTime, server_default=text("(datetime('now', 'utc'))"))


engine = create_engine(f'sqlite:///{SQLITE_PATH}')
Base.metadata.create_all(engine)


class MachineManager:
    def __init__(self):
        self.engine = engine
        self.Session = sessionmaker(bind=engine)

    def add_machine(self, machine_data):
        session = self.Session()
        machine = Machine(**machine_data)
        session.add(machine)
        session.commit()
        session.close()

    def add_machines(self, machines_data):
        session = self.Session()
        machines = [Machine(**data) for data in machines_data]
        session.bulk_save_objects(machines)
        session.commit()
        session.close()

    def delete_machine(self, machine_id):
        session = self.Session()
        machine = session.query(Machine).filter_by(id=machine_id).first()
        if machine:
            session.delete(machine)
            session.commit()
        session.close()

    def delete_machines(self, machine_ids):
        session = self.Session()
        machines = session.query(Machine).filter(Machine.id.in_(machine_ids)).all()
        for machine in machines:
            session.delete(machine)
        session.commit()
        session.close()

    def update_machine(self, machine_id, new_data):
        session = self.Session()
        machine = session.query(Machine).filter_by(id=machine_id).first()
        if machine:
            for key, value in new_data.items():
                setattr(machine, key, value)
            session.commit()
        session.close()

    def query_all_machines(self):
        session = self.Session()
        machines = session.query(Machine).all()
        session.close()
        return machines

    def query_machine_by_id(self, machine_id):
        session = self.Session()
        machine = session.query(Machine).filter_by(id=machine_id).first()
        session.close()
        return machine

    def query_machines_by_hostname(self, hostname):
        session = self.Session()
        machines = session.query(Machine).filter(Machine.hostname.like(f'%{hostname}%')).all()
        session.close()
        return machines

    def check_and_update_status(self):
        session = self.Session()
        machines = session.query(Machine).all()

        for machine in machines:
            current_time = datetime.datetime.utcnow()
            time_difference = current_time - machine.datetime

            if time_difference.total_seconds() > 1800:  # 超过30分钟（30分钟=1800秒）
                machine.status = False
        session.commit()
        session.close()

这里定义了一个主机表还有对应的操作类

ENV.py 是一个环境变量的文件

import os

from lib.readcfg import ReadConf

MAIN_PATH = os.path.abspath(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
SQLITE_PATH = os.path.join(MAIN_PATH, 'info.db')
LOG_PATH = os.path.join(MAIN_PATH, "agent.log")
AGENT_PATH = os.path.join(MAIN_PATH, "conf", "service.cfg")
AGENT_CFG = ReadConf(AGENT_PATH).return_info()

这里使用os.path 获取路径 不使用字符串众所周知使用字符串会导致不同平台不能兼容

readcfg.py 读取配置文件的类

from configobj import ConfigObj


class ReadConf:
    def __init__(self, CONF_PATH):
        self.config = ConfigObj(CONF_PATH, encoding="utf-8")

    def return_info(self):
        config_dict = {}
        for section, section_data in self.config.items():
            config_dict[section] = dict(section_data)
        return config_dict

返回一个字典数据

3.服务端主配置

service.cfg

[agent]
port = 12345 # socket 端口
host = 0.0.0.0 # socket 监听地址 

最后是服务端项目主文件

4.服务端主文件

server.py

import datetime
import json
import logging
import socket
import random
import pickle

from lib.ENV import LOG_PATH, AGENT_CFG
from lib.dataclass import MachineManager

logging.basicConfig(
    level=logging.DEBUG,  # 设置日志级别为 DEBUG，你可以根据需要选择不同的级别
    format="%(asctime)s [%(levelname)s] - %(message)s",
    datefmt='%Y-%m-%d %H:%M:%S',
    filename=LOG_PATH,  # 指定日志文件的路径
    filemode='a',  # 使用 'w' 模式以覆盖方式写入日志，使用 'a' 模式以追加方式写入日志
    encoding="utf-8"
)

# 生成两个随机质数
def generate_prime(bits):
    while True:
        num = random.getrandbits(bits)
        if is_prime(num):
            return num


# 检查一个数是否为质数
def is_prime(num):
    if num < 2:
        return False
    for i in range(2, int(num ** 0.5) + 1):
        if num % i == 0:
            return False
    return True


# 计算模反函数
def mod_inverse(a, m):
    for x in range(1, m):
        if (a * x) % m == 1:
            return x
    return None


# 生成密钥对
def generate_keypair(bits):
    p = generate_prime(bits)
    q = generate_prime(bits)
    n = p * q
    phi = (p - 1) * (q - 1)
    e = 65537  # 选择一个公钥指数
    d = mod_inverse(e, phi)  # 计算私钥指数
    return ((e, n), (d, n))


# 解密密文
def decrypt(private_key, ciphertext):
    d, n = private_key
    decrypted = [chr(pow(char, d, n)) for char in ciphertext]
    return "".join(decrypted)


def main():
    CONFIG = AGENT_CFG.get("agent", None)
    server_ip = CONFIG.get("host", None)
    server_port = int(CONFIG.get("port", None))
    bits = 10
    public_key, private_key = generate_keypair(bits)
    # 创建套接字并绑定到IP地址和端口
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.bind((server_ip, server_port))
    server_socket.listen(5)
    print(f"等待客户端连接在 {server_ip}:{server_port} 上...")
    machine_manager = MachineManager()
    while True:
        client_socket, client_address = server_socket.accept()
        logging.info(f"接收来自 {client_address} 的连接")
        client_socket.send(pickle.dumps(public_key))
        # 接收客户端信息
        data = b""  # 创建一个空的字节串
        while True:
            chunk = client_socket.recv(4096)
            if not chunk:
                break
            data += chunk
        client_info = json.loads(decrypt(private_key, pickle.loads(data)))  # 把传输过来的数据 使用私钥解密 再用json解析
        # print(client_info)
        logging.info(client_info)
        # print(client_info)
        client_info['datetime'] = datetime.datetime.strptime(client_info['datetime'], "%Y-%m-%d %I:%M:%S %p")
        # machine_manager.add_machine(client_info)
        # print(machine_manager.query_all_machines())
        if client_info:
            for i in machine_manager.query_all_machines():
                if i.hostname == client_info["hostname"]:
                    if i.cpu_info["Cores"] == client_info["cpu_info"]["Cores"] and i.cpu_info["Threads"] == client_info["cpu_info"]["Threads"]:
                        machine_manager.update_machine(i.id, client_info)
                if i.hostname != client_info["hostname"]:
                    if i.cpu_info["Cores"] != client_info["cpu_info"]["Cores"] and i.cpu_info["Threads"] != client_info["cpu_info"]["Threads"]:
                        machine_manager.add_machine(client_info)
            if not machine_manager.query_all_machines():
                machine_manager.add_machine(client_info)
        machine_manager.check_and_update_status()
        # 在这里可以处理客户端信息，保存到数据库或进行其他操作
        client_socket.close()


if __name__ == "__main__":
    main()

这里最后是判断新传入的主机是否为上一个主机 这个写的不太好希望大佬能指正

2.client端

接下来是client段的程序

client.cfg 客户端配置文件

[client]
id = 101   # 这里自定义 必须是Integer
service = 192.168.31.212  #这里是ip地址
service_port = 12345  # 这里是服务器接受端口

 

client.py

import datetime
import json
import pickle
import socket
import time
import psutil
import cpuinfo
from configobj import ConfigObj


class ReadConf:
    def __init__(self, CONF_PATH):
        self.config = ConfigObj(CONF_PATH, encoding="utf-8")

    def return_info(self):
        config_dict = {}
        for section, section_data in self.config.items():
            config_dict[section] = dict(section_data)
        return config_dict


# 加密明文
def encrypt(public_key, plaintext):
    e, n = public_key
    cipher = [pow(ord(char), e, n) for char in plaintext]
    return cipher


def get_uid_for_pid(pid):
    try:
        if is_windows():
            return psutil.Process(pid).username()
        else:
            return psutil.Process(pid).uids().real
    except psutil.NoSuchProcess:
        return None


def is_windows():
    import sys

    if sys.platform.startswith('win'):
        return True
    else:
        return False


def get_interface_info():
    interface_info = {}

    for interface, address in psutil.net_if_addrs().items():
        mac = None
        ipv4 = []
        ipv6 = []

        for addr in address:
            if is_windows():
                if addr.family == psutil.AF_LINK:  # MAC地址
                    mac = addr.address
                elif addr.family == socket.AF_INET:  # IPv4地址
                    ipv4.append(addr.address)

                elif addr.family == socket.AF_INET6:  # IPv6地址
                    ipv6.append(addr.address)
            else:
                if addr.family == socket.AF_PACKET:
                    mac = addr.address
                elif addr.family == socket.AF_INET:
                    ipv4.append(addr.address)
                elif addr.family == socket.AF_INET6:
                    ipv6.append(addr.address)
        interface_info[interface] = {
            "mac": mac,
            "ipv4": ipv4,
            "ipv6": ipv6
        }

    return interface_info


def get_client_info(host_id):
    # 获取主机名
    hostname = socket.gethostname()
    # 获取接口信息
    interfaces = get_interface_info()
    cpu_info_M_A_V: dict = cpuinfo.get_cpu_info()
    # 获取CPU使用率
    cpu_info = {
        "Cores": psutil.cpu_count(logical=False),  # 物理CPU核心数
        "Threads": psutil.cpu_count(logical=True),  # 逻辑CPU线程数）
        "CPU_Usage_%": [{cpu_index: cpu_usage} for cpu_index, cpu_usage in
                        enumerate(psutil.cpu_percent(interval=1, percpu=True), 0)],  # 每个CPU核心的使用情况
        "CPU_Model": cpu_info_M_A_V.get("brand_raw", None),
        "CPU_Architecture": cpu_info_M_A_V.get("arch", None),
        "CPU_Vendor": cpu_info_M_A_V.get("vendor_id", None),
    }

    # 获取内存信息
    memory_infos = psutil.virtual_memory()
    memory_info = {
        "Total_Memory_GB": round(memory_infos.total / (1024 ** 3), 2),
        "Available_Memory_GB": round(memory_infos.available / (1024 ** 3), 2),
        "Memory_Usage_%": round(memory_infos.percent, 2),
        "Used_Memory_GB": round(memory_infos.used / (1024 ** 3), 2),
        "Free_Memory_GB": round(memory_infos.free / (1024 ** 3), 2)
    }
    # print(memory_info)

    # 获取硬盘信息
    disk_info = psutil.disk_partitions()
    formatted_disk_info = []

    for partition in disk_info:
        partition_info = {
            "device": partition.device,
            "mountpoint": partition.mountpoint,
            "fstype": partition.fstype,
            "opts": partition.opts
        }
        formatted_disk_info.append(partition_info)

    # 获取进程信息
    processes = psutil.process_iter(attrs=['pid', 'name'])

    # 获取用户信息
    users = psutil.users()
    process_info = [{"pid": p.info['pid'], "server": p.info['name'], "uid": psutil.Process(p.info['pid']).username()}
                    for p in processes] if is_windows() else [
        {"pid": p.info['pid'], "server": p.info['name'], "uid": get_uid_for_pid(p.info['pid'])} for p in processes]
    disk_info = psutil.disk_partitions()
    disk_usage_info = []

    for partition in disk_info:
        usage = psutil.disk_usage(partition.mountpoint)
        disk_usage_info.append({
            "device": partition.device,
            "mountpoint": partition.mountpoint,
            "total_GB": round(usage.total / (1024 ** 3), 2),
            "used_GB": round(usage.used / (1024 ** 3), 2),
            "free_GB": round(usage.free / (1024 ** 3), 2),
            "percent_%": round(usage.percent, 2)
        })
    return {
        "id": host_id,
        "hostname": hostname,
        "cpu_info": cpu_info,
        "interfaces": interfaces,
        "memory_info": memory_info,
        "disk_info": formatted_disk_info,
        "disk_uses": disk_usage_info,
        "processes": process_info,
        "users": [{"user": u.name} for u in users],  # 获取用户ID (UID)
        "status": True,
        "datetime": datetime.datetime.now().strftime("%Y-%m-%d %I:%M:%S %p")
    }


def start():
    CONF = ReadConf("client.cfg").return_info().get("client")
    server_ip = CONF.get("service")
    # server_ip = "127.0.0.1"
    server_port = int(CONF.get("service_port"))
    ids = CONF.get("id")
    client_info = get_client_info(ids)
    # 创建套接字并连接到服务器
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.connect((server_ip, server_port))
    # 获取bits 的 tuple public——key
    public_key = pickle.loads(client_socket.recv(1024))
    # 发送客户端信息给服务器
    # print(client_info)
    client_socket.send(pickle.dumps(encrypt(public_key, json.dumps(client_info))))  # 使用公钥加密 发送数据
    # print(json.dumps(client_info))
    client_socket.close()


if __name__ == "__main__":
    while True:
        start()
        time.sleep(300)

这里获取的是 主机的  主机名 cpu信息 内存信息 硬盘信息 网卡信息 进程信息 用户信息

3.废话

效果图我还没有不过可以放加工页面显示通过发邮件实现我这里先放效果图，感兴趣的小伙伴可以耐心等等过段时间完善好在发出
 

个人观点，仅供参考