#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import socket
import psutil
import datetime
from openpyxl import Workbook
import xlrd2
from xlutils.copy import copy
import os


# 常量定义
system_name = '工作网邮件系统'

# 打印时间戳
host_name = socket.gethostname()
ipaddr = socket.gethostbyname(host_name)
current_time = datetime.datetime.now()
formatted_date = current_time.strftime("%Y-%m-%d")
filepath = "/var/log/system-check/" 
filename = filepath+"%s-%s.xlsx"%(formatted_date,ipaddr)
sheetname = formatted_date
formatted_time = current_time.strftime("%Y-%m-%d %H:%M:%S")
# print(formatted_time)


def table_head():
    table_head_list = [
        '巡检时间',
        '系统名称', '主机名', 'IP地址',
        'CPU核数', 'CPU使用率',
        '内存总量', '已用内存量', '剩余内存量', '内存使用率',
        'swap总量', '已用swap量', '剩余swap量', 'swap使用率',
        '硬盘名称', '硬盘位置', '硬盘总量', '硬盘已用量', '硬盘剩余量', '硬盘使用率',
    ]
    return table_head_list


def get_data():
    host_name = socket.gethostname()
    ipaddr = socket.gethostbyname(host_name)
    logical_cpus = psutil.cpu_count(logical=True)
    cpu_usage = psutil.cpu_percent(interval=1)
    memory = psutil.virtual_memory()
    mem_total = f"{memory.total / (1024 ** 3):.1f}GB"
    mem_used = f"{memory.used / (1024 ** 3):.1f}GB"
    mem_free = f"{memory.free / (1024 ** 3):.1f}GB"
    mem_usage = f"{memory.percent:.1f}%"
    swap_memory = psutil.swap_memory()
    swap_mem_total = f"{swap_memory.total / (1024 ** 3):.1f}GB"
    swap_mem_used = f"{swap_memory.used / (1024 ** 3):.1f}GB"
    swap_mem_free = f"{swap_memory.free / (1024 ** 3):.1f}GB"
    swap_mem_usage = f"{swap_memory.percent:.1f}%"
    data_list = [formatted_time, system_name, host_name, ipaddr, logical_cpus, cpu_usage, mem_total, mem_used, mem_free, mem_usage, swap_mem_total, swap_mem_used, swap_mem_free, swap_mem_usage]
    return data_list

def get_disk_data():
    disk = []
    disk_partitions = psutil.disk_partitions(all=False)
    for partition in disk_partitions:
        usage = psutil.disk_usage(partition.mountpoint)
        disk_name = partition.device
        disk_location = partition.mountpoint
        disk_total = f"{usage.total / (1024 ** 3):.1f}GB"
        disk_used = f"{usage.used / (1024 ** 3):.1f}GB"
        disk_free = f"{usage.free / (1024 ** 3):.1f}GB"
        disk_usage = f"{usage.percent:.1f}%"
        disk_info = [disk_name, disk_location, disk_total, disk_used, disk_free, disk_usage]
        disk.append(disk_info)
        # disk.extend(disk_info)
    return disk


def save_data_to_table(data, file_name, sheet_name):
    # 创建一个新的excel工作簿
    workbook = Workbook()
    # 创建一个新的工作表
    sheet = workbook.active
    sheet.title = sheet_name
    # 写入数据到Excel表格
    for row in data:
        sheet.append(row)
    # 保存Excel文件
    workbook.save(file_name)

def update_data(filename, value):
    index = len(value)  # 获取需要写入数据的行数
    workbook2 = xlrd2.open_workbook(filename)  # 打开工作簿
    sheets = workbook2.sheet_names()  # 获取工作簿中的所有表格
    worksheet = workbook2.sheet_by_name(sheets[0])  # 获取工作簿中所有表格中的的第一个表格
    rows_old = worksheet.nrows  # 获取表格中已存在的数据的行数
    # print(rows_old)
    new_workbook = copy(workbook2)  # 将xlrd对象拷贝转化为xlwt对象
    # print(new_workbook)
    new_worksheet = new_workbook.get_sheet(0)  # 获取转化后工作簿中的第一个表格
    for ii in range(0, index):
        for jj in range(0, len(value[ii])):
            new_worksheet.write(ii+rows_old, jj, value[ii][jj])  # 追加写入数据,注意是从i+rows_old行开始写入
            # print(value[ii][jj])
    new_workbook.save(filename)  # 保存工作簿


def file_exists(filename):
    return os.path.exists(filename)


def main():
    all_data = []

    head_data = table_head()
    all_data.append(head_data)

    table_data = get_data()

    disk_data = get_disk_data()
    for i in disk_data:
        table_data.extend(i)
        break
    all_data.append(table_data)

    disk_second_data = []
    for i in disk_data[1:]:
        second = []
        for j in get_data():
            j = ''
            second.append(j)
        disk_second_data.append(second + i)
        all_data.append(second + i)

    # print(all_data[0])
    # print(all_data[1:-1])

    # save_data_to_table(all_data, filename, sheetname)

    if file_exists(filename):
        update_data(filename, all_data[1:])
    else:
        save_data_to_table([all_data[0]], filename, sheetname)
        update_data(filename, all_data[1:])


if __name__ == '__main__':
    main()