nohang/os04

#!/usr/bin/env python3
"""
sort processes by oom_score
"""

# REFACTORING IN PROGRESS

from operator import itemgetter
from os import listdir
from argparse import ArgumentParser


"""#######################################################################79"""

# define funtcions


def pid_to_cmdline(pid):
    """
    Get process cmdline by pid.

    pid: str pid of required process
    returns string cmdline
    """
    with open('/proc/' + pid + '/cmdline') as f:
        return f.read().replace('\x00', ' ').rstrip()


def pid_to_name(pid):
    """
    Get process name by pid.

    pid: str pid of required process
    returns string process_name
    """
    try:
        with open('/proc/' + pid + '/status') as f:
            f.seek(6)
            for line in f:
                return line[:-1]
    except FileNotFoundError:
        return ''
    except ProcessLookupError:
        return ''
    except UnicodeDecodeError:
        with open('/proc/' + pid + '/status', 'rb') as f:
            f.seek(6)
            return f.read(15).decode(
                'utf-8', 'ignore').partition('\n')[0]



def kib_to_mib(num):
    """Convert KiB values to MiB values."""
    return round(num / 1024.0)

def human(num):
    '''KiB to MiB'''
    return str(round(num / 1024.0)).rjust(6, ' ')


def get_max_pid_len():
    with open('/proc/sys/kernel/pid_max') as file:
        for line in file:
            return len(line.strip())


def get_max_vm_rss_len():
    pass



def pid_to_status_units(pid):

    with open('/proc/' + pid + '/status', 'rb') as f:
        f_list = f.read().decode('utf-8', 'ignore').split('\n')

        for i in range(len(f_list)):

            if i is 1:
                name = f_list[0].split('\t')[1]
                print(i, 'Name:', name)

            if i is uid_index:
                uid = f_list[i].split('\t')[2]
                print(i, 'UID:', uid)

            if i is vm_rss_index:
                pass
                vm_rss = kib_to_mib(int(f_list[i].split('\t')[1][:-3]))
                print(i, 'rss:', vm_rss)

            if i is vm_swap_index:
                pass
                vm_swap = kib_to_mib(int(f_list[i].split('\t')[1][:-3]))
                print(i, 'swap:', vm_swap)




"""#######################################################################79"""


# parse input

# todo: input validation


parser = ArgumentParser()

parser.add_argument(
    '--num',
    '-n',
    help="""max number of lines; default: 99999""",
    default=None,
    type=str
)

parser.add_argument(
    '--len',
    '-l',
    help="""max cmdline length; default: 99999""",
    default=None,
    type=str
)

args = parser.parse_args()

display_cmdline = args.len
num_lines = args.num
if num_lines is None:
    num_lines = 20

if display_cmdline is None:
    display_cmdline = 20


"""#######################################################################79"""

# find VmRSS, VmSwap and UID positions in /proc/*/status

with open('/proc/self/status') as file:
    status_list = file.readlines()

# список имен из /proc/*/status для дальнейшего поиска позиций VmRSS and VmSwap
status_names = []
for s in status_list:
    status_names.append(s.split(':')[0])

vm_rss_index = status_names.index('VmRSS')
vm_swap_index = status_names.index('VmSwap')
uid_index = status_names.index('Uid')















"""#######################################################################79"""

# get sorted list with pid, oom_score, oom_score_adj, cmdline

# здесь же искать вообще все.


oom_list = []

for pid in listdir('/proc'):

    # пропускаем элементы, состоящие не из цифр и PID 1
    if pid.isdigit() is False or pid == '1':
        continue


    '''


    name = pid_to_name(pid)
    '''

    try:
        cmdline = pid_to_cmdline(pid)
        if cmdline == '':
            continue

        with open('/proc/' + pid + '/oom_score') as file:
            oom_score = int(file.readlines()[0][:-1])

        with open('/proc/' + pid + '/oom_score_adj') as file:
            oom_score_adj = int(file.readlines()[0][:-1])

    except FileNotFoundError:
        continue

    except ProcessLookupError:
        continue

    oom_list.append((pid, oom_score, oom_score_adj, cmdline))

# list sorted by oom_score
oom_list_sorted = sorted(oom_list, key=itemgetter(1), reverse=True)





"""#######################################################################79"""


# print table


max_pid_len = get_max_pid_len()

if display_cmdline == '0':

    print(
        'oom_score oom_score_adj   UID{}PID Name             VmRSS   VmSwap'.format(' ' * (max_pid_len - 2)))

    print('--------- ------------- ----- {} --------------- -'
          '------- --------'.format('-' * max_pid_len))

else:

    print('oom_score oom_score_adj   UID{}PID Name          '
          '   VmRSS   VmSwap   cmdline'.format(' ' * (max_pid_len - 2)))

    print('--------- ------------- ----- {} --------------- -'
          '------- -------- -------'.format('-' * max_pid_len))

# итерируемся по сортированному списку oom_score, печатая name, pid etc


for i in oom_list_sorted[:int(num_lines)]:

    pid = i[0]
    oom_score = i[1]
    oom_score_adj = i[2]
    cmdline = i[3]

    try:
        # читать часть файла не дальше VmSwap - когда-нибудь
        with open('/proc/' + pid + '/status') as file:
            status_list = file.readlines()

        vm_rss = int(status_list[vm_rss_index].split(':')[1].split(' ')[-2])
        vm_swap = int(status_list[vm_swap_index].split(':')[1].split(' ')[-2])
        name = status_list[0][:-1].split('\t')[1]
        uid = status_list[uid_index].split('\t')[1]

    except FileNotFoundError:
        continue

    except ProcessLookupError:
        continue

    print(
        '{} {} {} {} {} {} M {} M {}'.format(
            str(oom_score).rjust(9),
            str(oom_score_adj).rjust(13),
            uid.rjust(5),
            str(pid).rjust(max_pid_len),
            name.ljust(15),
            human(vm_rss),
            human(vm_swap),
            cmdline[:int(display_cmdline)]
        )
    )