#coding: utf-8
# +-------------------------------------------------------------------
# | 宝塔Linux面板
# +-------------------------------------------------------------------
# | Copyright (c) 2015-2099 宝塔软件(http://bt.cn) All rights reserved.
# +-------------------------------------------------------------------
# | Author: hwliang <hwl@bt.cn>
# +-------------------------------------------------------------------

#--------------------------------
# 进程监控模块
#--------------------------------
from psutil import cpu_count,pids,Process,cpu_times
from json import dumps
import os,sys

os.chdir(os.getenv('BT_PANEL'))
sys.path.insert(0, "/")

import copy, threading,struct,time,db
from cachelib import SimpleCache


class process_task:
    __pids = []
    __last_times = {}
    __last_dates = {}
    __write_last = {}
    __write_dates = {}
    __read_last = {}
    __read_dates = {}
    __cpu_count = cpu_count()
    __cache = SimpleCache(5000)
    last_net_process = {}
    last_net_process_time = 0
    __process_net_list = {}
    __process_object = {}
    __insert_time = 0
    __last_cpu_time = 0
    old_key = 'old_cpu_times'
    new_key = 'new_cpu_times'
    new_info = {}
    old_info = {}


    def __init__(self):
        _sql = db.Sql().dbfile('system')
        csql = '''CREATE TABLE IF NOT EXISTS `process_top_list` (
  `id` INTEGER PRIMARY KEY AUTOINCREMENT,
  `cpu_top` REAL,
  `memory_top` REAL,
  `disk_top` REAL,
  `net_top` REAL,
  `all_top` REAL,
  `addtime` INTEGER
)'''
        _sql.execute(csql, ())
        _sql.close()

    def get_pids(self):
        '''
            @name 获取pid列表
            @author hwliang<2021-09-04>
            @return None
        '''
        self.__pids = pids()

    def get_cpu_time(self):
        s = cpu_times()

        return s.user + s.system  + s.idle

    def get_old(self):
        if self.old_info: return True
        data = self.__cache.get(self.old_key)
        if not data: return False
        if not data: return False
        self.old_info = data
        del(data)
        return True


    # def get_cpu_percent(self,pid,cpu_times,cpu_time):
    #     self.get_old()
    #     percent = 0.00
    #     process_cpu_time = self.get_process_cpu_time(cpu_times)
    #     if not self.old_info: self.old_info = {}
    #     if not pid in self.old_info:
    #         self.new_info[pid] = {}
    #         self.new_info[pid]['cpu_time'] = process_cpu_time
    #         return percent
    #     percent = round(100.00 * (process_cpu_time - self.old_info[pid]['cpu_time']) / (cpu_time - self.old_info['cpu_time']),2)
    #     self.new_info[pid] = {}
    #     self.new_info[pid]['cpu_time'] = process_cpu_time
    #     if percent > 0: return percent
    #     return 0.00

    def get_process_cpu_time(self,cpu_times):
        cpu_time = 0.00
        for s in cpu_times: cpu_time += s
        return cpu_time

    def get_cpu_percent(self,pid,cpu_time_total,s_cpu_times):
        '''
            @name 获取pid的cpu占用率
            @author hwliang<2021-09-04>
            @param pid 进程id
            @param cpu_time_total 进程总cpu时间
            @return 占用cpu百分比
        '''
        stime = time.time()
        if pid in self.__last_times:
            old_time = self.__last_times[pid]
        else:
            self.__last_times[pid] = cpu_time_total
            self.__last_dates[pid] = stime
            return 0

        cpu_percent = round(100.00 * float(cpu_time_total - old_time) / (s_cpu_times - self.__last_cpu_time),2)
        self.__last_times[pid] = cpu_time_total
        self.__last_dates[pid] = stime
        if cpu_percent > 100: cpu_percent = 99
        return cpu_percent


    def get_io_write(self,pid,io_write):
        disk_io_write = 0
        stime = time.time()
        if pid in self.__write_last:
            old_write = self.__write_last[pid]
        else:
            self.__write_last[pid] = io_write
            self.__write_dates[pid] = stime
            return disk_io_write

        io_end = (io_write - old_write)
        if io_end > 0:
            disk_io_write = int(io_end / (stime - self.__write_dates[pid]))

        self.__write_last[pid] = io_write
        self.__write_dates[pid] = stime
        return disk_io_write


    def get_io_read(self,pid,io_read):
        disk_io_read = 0
        stime = time.time()
        if pid in self.__read_last:
            old_read = self.__read_last[pid]
        else:
            self.__read_last[pid] = io_read
            self.__read_dates[pid] = stime
            return disk_io_read

        io_end = (io_read - old_read)
        if io_end > 0:
            disk_io_read = int(io_end / (stime - self.__read_dates[pid]))

        self.__read_last[pid] = io_read
        self.__read_dates[pid] = stime
        return disk_io_read

    def read_file(self,filename):
        f = open(filename,'rb')
        result = f.read()
        f.close()
        return result.decode().replace("\u0000"," ").strip()




    def get_process_net_list(self):
        w_file = '/dev/shm/bt_net_process'
        if not os.path.exists(w_file): return
        self.last_net_process = self.__cache.get('net_process')
        self.last_net_process_time = self.__cache.get('last_net_process')
        net_process_body  = self.read_file(w_file)
        if not net_process_body: return
        net_process = net_process_body.split('\n')
        for np in net_process:
            if not np: continue
            tmp = {}
            np_list = np.split()
            if len(np_list) < 5: continue
            tmp['pid'] = int(np_list[0])
            tmp['down'] = int(np_list[1])
            tmp['up'] = int(np_list[2])
            tmp['down_package'] = int(np_list[3])
            tmp['up_package'] = int(np_list[4])
            self.__process_net_list[str(tmp['pid'])] = tmp
        self.__cache.set('net_process',self.__process_net_list,600)
        self.__cache.set('last_net_process',time.time(),600)

    def get_process_network(self,pid):
        '''
            @name 获取进程网络流量
            @author hwliang<2021-09-13>
            @param pid<int> 进程ID
            @return tuple
        '''

        if not self.__process_net_list:
            self.get_process_net_list()
        if not self.last_net_process_time:
            return 0,0,0,0
        if not pid in self.__process_net_list.keys():
            return 0,0,0,0

        if not pid in self.last_net_process:
            return self.__process_net_list[pid]['up'],self.__process_net_list[pid]['up_package'],self.__process_net_list[pid]['down'],self.__process_net_list[pid]['down_package']

        up = int((self.__process_net_list[pid]['up'] - self.last_net_process[pid]['up']) / (time.time() - self.last_net_process_time))
        down = int((self.__process_net_list[pid]['down'] - self.last_net_process[pid]['down']) / (time.time() - self.last_net_process_time))
        up_package = int((self.__process_net_list[pid]['up_package'] - self.last_net_process[pid]['up_package']) / (time.time() - self.last_net_process_time))
        down_package = int((self.__process_net_list[pid]['down_package'] - self.last_net_process[pid]['down_package']) / (time.time() - self.last_net_process_time))
        if up < 0: up = 0
        if down < 0: down = 0
        if up_package < 0: up_package = 0
        if down_package < 0: down_package = 0
        return up,up_package,down,down_package


    def get_process_username(self,pid):
        '''
            @name 获取进程用户名
            @param pid 进程id
            @return 用户名
        '''
        try:
            import pwd
            return pwd.getpwuid(os.stat('/proc/' + str(pid)).st_uid).pw_name
        except:
            return 'root'


    def get_monitor_list(self,stime = None):
        '''
            @name 获取监控列表
            @author hwliang<2021-09-04>
            @return list
        '''
        self.get_pids()
        process_info_list = []
        total_cpu_precent = 0.0
        my_pid = os.getpid()
        if type(self.new_info) != dict: self.new_info = {}
        all_cpu_time = self.get_cpu_time()
        self.new_info['time'] = time.time()

        for pid in self.__pids:
            try:
                if pid < 100: continue
                if pid == my_pid: continue
                if not pid in self.__process_object.keys():
                    self.__process_object[pid] = Process(pid)
                try:
                    if self.__process_object[pid].status() == 'terminated':
                        self.__process_object[pid] = Process(pid).create_time
                except:
                    self.__process_object[pid] = Process(pid)
                p = self.__process_object[pid]

                if p.name() in ['System Idle Process']:
                    continue

                process_info = {}
                process_info['cpu_percent'] = self.get_cpu_percent(str(pid),sum(p.cpu_times()),all_cpu_time) #self.get_cpu_percent(pid,int(sum(p.cpu_times())),self.new_info['cpu_time'])         # CPU使用率
                total_cpu_precent += process_info['cpu_percent']
                process_info['memory'] = p.memory_info().rss                                            # 内存占用
                if not process_info['memory']: continue

                io_counters = p.io_counters()
                process_info['disk_read'] = self.get_io_read(pid,io_counters.read_bytes)                # 读取磁盘字节数
                process_info['disk_write'] = self.get_io_write(pid,io_counters.write_bytes)             # 写入磁盘字节数
                process_info['disk_total'] = process_info['disk_read'] + process_info['disk_write']     # 磁盘总读写

                process_info['up'],process_info['up_package'],process_info['down'],process_info['down_package'] = self.get_process_network(str(pid))

                process_info['net_total'] = process_info['up'] + process_info['down']                   # 网络总流量
                process_info['package_total'] = process_info['up_package'] + process_info['down_package'] # 网络总包数

                if not process_info['cpu_percent'] and not process_info['disk_total'] and not process_info['net_total']: continue

                process_info['pid'] = pid
                process_info['name'] = p.name()
                process_info['cmdline'] = p.cmdline()                                  # 启动命令
                process_info['create_time'] = int(p.create_time())                                                          # 创建时间
                process_info['connect_count'] = len(p.connections())                                                      # 连接数
                process_info['username'] = p.username()                                    # 进程用户名
                process_info_list.append(process_info)
            except:
                continue

        self.__last_cpu_time = all_cpu_time
        self.__process_net_list.clear()
        self.insert_db(process_info_list,stime)
        # import public
        # for pp in sorted(process_info_list,key=lambda x:x['cpu_percent'],reverse=True):
        #     public.print_log("name: {}, cpu_percent: {}".format(pp['name'], pp['cpu_percent']))
        if total_cpu_precent > 100: total_cpu_precent = 100
        return total_cpu_precent

    def get_expire_time(self):
        '''
            @name 获取过期时间
            @return int
        '''
        filename = 'data/control.conf'
        _day = 30
        if os.path.exists(filename):
            try:
                conf = self.read_file(filename)
                if conf: _day = int(conf)
            except: pass
        return _day * 86400



    def insert_db(self,process_info_list,_time):
        '''
            @name 插入数据库
            @param process_info_list list
            @return bool
        '''
        if not process_info_list: return
        all_top,cpu_top,disk_top,net_top,memory_top = self.get_top_list(process_info_list)

        with db.Sql().dbfile('system') as _sql:
            if not _time:
                _time = int(time.time())
            _sql.table('process_top_list').insert({
                'all_top':dumps(all_top),
                'cpu_top':dumps(cpu_top),
                'disk_top':dumps(disk_top),
                'net_top':dumps(net_top),
                'memory_top':dumps(memory_top),
                'addtime':_time
            })

            # 删除过期数据
            if self.__insert_time and _time - self.__insert_time > 3600:
                self.__insert_time = _time
                _sql.table('process_top_list').where('addtime<?',self.get_expire_time()).delete()
            _sql.close()



    def get_top_list(sekf,process_info_list):
        '''
            @name 排序
            @param process_info_list list
            @return list
        '''
        process_info_list = sorted(process_info_list,key=lambda x:[x['cpu_percent'],x['disk_total'],x['net_total'],x['memory']],reverse=True)
        top_num = 5
        all_top = []
        for p in process_info_list[:top_num]:
            _line = [p['cpu_percent'],p['disk_read'],p['disk_write'],p['memory'],p['up'],p['down'],p['pid'],p['name'],p['cmdline'],p['username'],p['create_time']]
            all_top.append(_line)


        process_info_list = sorted(process_info_list,key=lambda x:x['cpu_percent'],reverse=True)
        cpu_top = []
        for p in process_info_list[:top_num]:
            if not p['cpu_percent']: continue
            _line = [p['cpu_percent'],p['pid'],p['name'],p['cmdline'],p['username'],p['create_time']]
            cpu_top.append(_line)

        process_info_list = sorted(process_info_list,key=lambda x:x['disk_total'],reverse=True)
        disk_top = []
        for p in process_info_list[:top_num]:
            if not p['disk_total']: continue
            _line = [p['disk_total'],p['disk_read'],p['disk_write'],p['pid'],p['name'],p['cmdline'],p['username'],p['create_time']]
            disk_top.append(_line)

        process_info_list = sorted(process_info_list,key=lambda x:x['net_total'],reverse=True)
        net_top = []
        for p in process_info_list[:top_num]:
            if not p['net_total']: continue
            _line = [p['net_total'],p['up'],p['down'],p['connect_count'],p['package_total'],p['pid'],p['name'],p['cmdline'],p['username'],p['create_time']]
            net_top.append(_line)

        process_info_list = sorted(process_info_list,key=lambda x:x['memory'],reverse=True)
        memory_top = []
        for p in process_info_list[:top_num]:
            if not p['memory']: continue
            _line = [p['memory'],p['pid'],p['name'],p['cmdline'],p['username'],p['create_time']]
            memory_top.append(_line)

        return all_top,cpu_top,disk_top,net_top,memory_top




if __name__ == '__main__':
    # net = process_network_total()
    # threading.Thread(target = net.start,args=()).start()

    p = process_task()
    while True:
        p.get_monitor_list()
        time.sleep(1)
        print("-"*50)
