from subprocess import check_output
from time import sleep
from prometheus_client import start_http_server, Gauge
import argparse


def get_cmd_window_info():
    cmd_output = check_output('tasklist /v /fo csv').decode('GBK')

    lines = cmd_output.split('\n')[1:-1]  # Skip header and empty line
    cmd_info = []
    for line in lines:
        fields = line.strip().split('","')
        image = fields[0][1:]
        pid = int(fields[1])
        mem_use = int(fields[4][:-2].replace(',', '')) * 1024
        title = fields[8].split(':')[-1].strip('"').strip()

        cmd_info.append({'image': image, 'pid': pid, 'mem_use': mem_use, 'title': title})
    return cmd_info


if __name__ == '__main__':
    print()
    print('*' * 20, "Windows process exporter For Prometheus", '*' * 20)
    print("You can directly start to use default port 8760,\nor add --port port_num as a start args to use port_num")
    # 创建一个Gauge指标
    window_memory_used_by_process = Gauge('window_memory_used_by_process', 'Window process Memory Used',
                                          ['title', 'pid', 'image'])

    # 解析命令行参数
    parser = argparse.ArgumentParser()
    parser.add_argument('--port', type=int, default=8760, help='Port number to listen on')
    args = parser.parse_args()

    # 启动HTTP服务器，监听指定的端口，并添加基本身份验证
    start_http_server(args.port)

    while True:
        # 获取cmd.exe程序启动的命令行窗口信息

        cmd_window_info = get_cmd_window_info()

        # 更新指标的值
        window_memory_used_by_process._metrics.clear()
        for info in cmd_window_info:
            window_memory_used_by_process.labels(title=info['title'], pid=str(info['pid']),
                                                 image=str(info['image'])).set(
                info['mem_use'])

        # 每隔30秒更新一次指标的值
        sleep(30)
