#!/usr/bin/env python
# -*- coding:utf-8 -*-
import psutil

# 获取当前系统的登录用户数
login_users = len(psutil.users())


def bytes2human(n):
    """
    字节转换工具
    http://code.activestate.com/recipes/578019
    #bytes2human(10000)
    '9.8K'
    #>>> bytes2human(100001221)
    '95.4M'
    :param n:
    :return:
    """

    symbols = ('K', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y')
    prefix = {}
    for i, s in enumerate(symbols):
       # 0  K
       # 1  M
        prefix[s] = 1 << (i + 1) * 10
       # {"K": 1024,
       #  "M": 1048576,
       #  "Y": ... }
    for s in reversed(symbols):
        if n >= prefix[s]:
            value = float(n) / prefix[s]
            return '{:.1f}{}'.format(value, s)
    return "%sB" % n


# 获取 CPU 信息
def get_cpu():
    # cpu 实例数
    cpu_count = len(psutil.Process().cpu_affinity())

    # 每颗 cpu 的物理核心数
    cpu_py_count = psutil.cpu_count(logical=False)

    # cpu 信息
    cpu_info = psutil.cpu_times_percent(interval=0.1)

    return cpu_count, cpu_py_count, cpu_info


# 打印 CPU 信息
def display_cpu():
    cpu_temple = """
************************** CPU 信息 **************************
CPU 颗数： {cpu_count}       CPU 总使用率: {cpu_time}%
物理核心数: {cpu_py_count}    IO 等待 CPU 空闲率: {io_wait_time}%  
"""

    cpu_count, cpu_py_count, cpu_info = get_cpu()
    import sys
    plf = sys.platform
    print(cpu_temple.format(
        cpu_count=cpu_count,
        cpu_py_count=cpu_py_count,
        cpu_time=cpu_info.user + cpu_info.system,

        # iowait Mac 系统不支持,在 mac 系统下需要把 io_wait_time 赋值为空
        # io_wait_time=""

        io_wait_time=cpu_info.iowait if plf == 'linux' else ''
    ))


def get_mem():
    # 内存信息
    return psutil.virtual_memory()

# mem_info = psutil.virtual_memory()

def display_meme():
    mem_info = get_mem()
    print("*" * 16, "Disk 信息", "*" * 16)
    fileds = ['total', 'free', 'used']
    for name in mem_info._fields:
        # if name != 'percent':
        if name in fileds:
            value = getattr(mem_info, name)
            value = bytes2human(value)
            print("{:<10s} : {}".format(name.capitalize(), value))


# 获取 DISK 信息
def get_disk_part():
    return psutil.disk_partitions(all=False)


def display_disk():
    disk_fields = ("Device", "Total", "Used", "Free", "Use ", "Type", "Mount")
    temple_title = '{:<23s} {:<8} {:<8} {:<8} {:<3s}% {:>8s}  {:<s}'
    temple = '{:<23s} {:<8s} {:<8s} {:<8s} {:3.1f}% {:>8s}  {:<s}'

    disk_part = get_disk_part()
    disk_detail = temple_title.format(*disk_fields)

    print("*" * 35, "Disk 信息", "*" * 35)
    print(disk_detail)
    for part in disk_part:
        usage = psutil.disk_usage(part.mountpoint)  # 获取磁盘使用率
        """
        注意: UNIX通常为root用户保留总磁盘空间的5％。
             UNIX上的 Total 和 Used 字段是指总体使用和已用空间，
             而 free 表示用户可用空间， 百分比表示用户使用率（请参阅 源代码）。
             这就是为什么百分比值可能比你期望的要高5％。
        """
        print(temple.format(
            part.device,
            bytes2human(usage.total),
            bytes2human(usage.used),
            bytes2human(usage.free),
            usage.percent,
            part.fstype,
            part.mountpoint))


# 获取网络信息
def get_net_info():
    """获取到网络的基础信息和地址"""
    # 导入 socket 模块，用于识别 ipv4 地址信息
    import socket

    # 定义一个空字典，用于存放获取到的数据
    net_info_dic = {}

    # 获取网卡的状态信息，返回的是个字典
    if_stats_dic = psutil.net_if_stats()

    # 获取网卡的地址信息，返回的是个字典
    if_addrs_dic = psutil.net_if_addrs()

    # 获取网卡的 io 信息，返回的是个字典
    io_counters_obj = psutil.net_io_counters()

    for nic, addrs in if_addrs_dic.items():
        if not nic.startswith('lo'):
            # 首先得到网卡的状态，并把数据填充到字典中
            net_info_dic[nic] = {'nic_stat': if_stats_dic[nic].isup}
            # 示例：net_info_dic = {'eth0': {'nic_stat': True}

            for addr in addrs:  # 循环每块网卡地址的每类信息

                if addr.family == socket.AF_INET:  # 我们只需要 ipv4 的信息
                    # 更新相关信息到每个网卡的字典中
                    net_info_dic[nic].update({
                        'ip': addr.address,
                        'netmask': addr.netmask
                     })

    # 假如某个网卡不正常，它的 ip 地址信息可能不存在，
    # 这样的话，这个网卡的字典信息中将会缺失 ip 和 netmask 的键；
    # 为了后期取数方便，所以在此处进行处理，
    # 把这两个键添加到字典中，并且赋值为空字符串

    for item in net_info_dic.values():
        item.setdefault('ip', '')
        item.setdefault('netmask', '')


    # 更新网卡 io 信息到字典中
    net_info_dic['io_info'] = {
        # 发送字节数(bytes),除以 1024 得到 kb
        'bytes_sent': bytes2human(io_counters_obj.bytes_sent),
        'bytes_recv': bytes2human(io_counters_obj.bytes_recv),  # 接收字节数
        'packe_sent': io_counters_obj.packets_sent,  # 发送数据包数
        'packe_recv': io_counters_obj.packets_recv}

    return net_info_dic

def display_net():
    """
    获取到网卡格式化后的信息，并格式化打印到屏幕上
    :return:
    """
    net_info_dic = get_net_info()
    print(net_info_dic)
    temple_addr = """
网卡名: {nic}  状态: {stat}
     IP: {ip} 掩码: {net_mask}
      """
    temple_io = """
    
当前系统网络 IO 信息(统计的是所有网卡的信息)
      当前发送字节数: {bytes_sent}  当前接收字节数: {bytes_recv}
    """

    print("*" * 30, "Network 信息", "*" * 30)
    print(temple_io.format(bytes_sent=net_info_dic['io_info'].get('bytes_sent', ''),
                           bytes_recv=net_info_dic['io_info'].get('bytes_recv', '')
                          )
    )

    net_info_dic.pop('io_info')

    for nic, item in net_info_dic.items():
        print(temple_addr.format(nic=nic,
                                 stat=item['nic_stat'],
                                 ip=item.get('ip', ''),
                                 net_mask=item.get('netmask', '')
                                 )
              )


def man():
    import sys
    display_dic = {
        "1": {"title": "查看 CPU 信息", "func": display_cpu},
        "2": {"title": "查看 Memory 信息", "func": display_meme},
        "3": {"title": "查看 Disk 信息", "func": display_disk},
        "4": {"title": "查看 Network 信息", "func": display_net}
    }
    print("{}Python系统性能信息收集".format('*' * 10))  # 首先打印出菜单标题
    for k, item in display_dic.items():
        print(k, item['title'])
    while True:
        inp = input('请选择>>:')
        if inp == 'q':
            sys.exit("系统退出")
        if inp in display_dic.keys():
            # python3.6+
            print(f"{'*' * 10}Python 系统性能信息收集程序")         
            # python3.x
            print("{}Python 系统性能信息收集程序".format('*' * 10))          
            for k, item in display_dic.items():
                print(k, item['title'])
            print()
            exec_func = display_dic[inp]['func']
            exec_func()
            print()

if __name__ == "__main__":
    man()


###########################

"""
{'eth0': snetio(bytes_sent=101643418, bytes_recv=86163253, packets_sent=286285, packets_recv=326059, errin=0, errout=0, dropin=0, dropout=0),
 'lo': snetio(bytes_sent=224151, bytes_recv=224151, packets_sent=3347, packets_recv=3347, errin=0, errout=0, dropin=0, dropout=0)}
"""

"""
psutil.net_if_addrs（）
将与安装在系统上的每个NIC（网络接口卡）关联的地址作为字典返回，该字典的关键字是NIC名称，值是分配给NIC的每个地址的命名元组列表。每个命名的元组包括5个字段：

系列：地址系列， AF_INET， AF_INET6 或psutil.AF_LINK，指MAC地址。
地址：主NIC地址（始终设置）。
网络掩码：网络掩码地址（可能是None）。
广播：广播地址（可能None）。
ptp：代表“点对点”; 它是点对点接口（通常是VPN）的目标地址。广播和PT是互斥的。可能是None。
"""

"""
psutil.net_if_stats（）
{'eth0': snicstats(isup=True, duplex=<NicDuplex.NIC_DUPLEX_FULL: 2>, speed=1000, mtu=1500),
 'lo': snicstats(isup=True, duplex=<NicDuplex.NIC_DUPLEX_UNKNOWN: 0>, speed=0, mtu=65536)}
将安装在系统上的每个NIC（网络接口卡）的信息作为字典（其关键字是NIC名称和值）返回，其中包含以下字段的命名元组：

isup：表示NIC是否启动并运行的布尔值。  这里我们只关心 是否是启动状态
duplex：双工通信类型; 它可以是NIC_DUPLEX_FULL，NIC_DUPLEX_HALF或者 NIC_DUPLEX_UNKNOWN。
speed：以兆位（MB）表示的NIC速度，如果无法确定（例如'localhost'），则将其设置为0。
mtu：NIC的最大传输单位，以字节表示。
"""

"""
我们经常把socket翻译为套接字，socket是在应用层和传输层之间的一个抽象层，
它把TCP/IP层复杂的操作抽象为几个简单的接口供应用层调用已实现进程在网络中通信
python 的 socket 包中有 网卡接口的类型，利用这个可以实现一个映射关系

socket.AF_INET = <AddressFamily.AF_INET: 2>
psutil.AF_LINK = <AddressFamily.AF_LINK: 17>


"""