#!/usr/bin/env python3
"""
Fork bomb detector and deactvator
:(){ :|:& };:
while true; do bash -c ':(){ :|:& };: &'; done
perl -MPOSIX -e '$pid=fork; exit if $pid; setid; fork while(1)'
while true; do (sleep 60 &); done
perl -e "fork while fork" &
echo '. ./_&. ./_'>_;. ./_
"""
import os
from sys import stdout
from operator import itemgetter
from time import sleep, time
from signal import SIGKILL, SIGSTOP, SIGCONT

self_pid = str(os.getpid())

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

status_names = []

for s in status_list:
    status_names.append(s.split(':')[0])

ppid_index = status_names.index('PPid')
uid_index = status_names.index('Uid')

PROC_NUM = 1000
TIME = 5
PRINT_NPROC = 10

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

# FUNCS

def rline1(path):
    """read 1st line from path."""
    try:
        with open(path) as f:
            for line in f:
                return line[:-1]
    except FileNotFoundError:
        return ''
    except ProcessLookupError:
        return ''


def pid_to_uid(pid):
    '''return euid'''
    try:

        with open('/proc/' + pid + '/status') as f:
            for n, line in enumerate(f):
                if n is uid_index:
                    return line.split('\t')[2]
    except FileNotFoundError:
        return ''
    except ProcessLookupError:
        return ''


def pid_to_name(pid):
    try:
        with open('/proc/' + pid + '/status') as f:
            for line in f:
                return line[:-1].split('\t')[1]
    except FileNotFoundError:
        return ''
    except ProcessLookupError:
        return ''


def pid_to_ppid(pid):
    try:
        with open('/proc/' + pid + '/status') as f:
            for n, line in enumerate(f):
                if n is ppid_index:
                    return line.split('\t')[1].strip()
    except FileNotFoundError:
        return ''
    except ProcessLookupError:
        return ''


def pid_to_cmdline(pid):
    try:
        with open('/proc/' + pid + '/cmdline') as f:
            try:
                for line in f:
                    return line.replace('\x00', ' ').strip()
            except IndexError:
                return ''
    except FileNotFoundError:
        return ''
    except ProcessLookupError:
        return ''


def pid_to_state(pid):
    try:
        return rline1('/proc/' + pid + '/stat').rpartition(')')[2][1]
    except FileNotFoundError:
        return ''
    except ProcessLookupError:
        return ''
    except IndexError:
        return ''


def find_bomb_pattern(pid_set):
    # need rename to find_bomb_name
    print('Find name of bomb processes')
    t0 = time()
    name_dict = dict()
    for pid in pid_set:
        name = pid_to_name(pid)
        if name == '':
            continue
        if name not in name_dict:
            name_dict[name] = 1
        else:
            name_dict[name] += 1
    x = name_dict.items()
    y = sorted(x, key=itemgetter(1))
    print(y[::-1][:9])
    print('Time:', round(time() - t0, 3))
    return y[-1][0]


def find_bomb_cmdline():
    print('Find bomb cmdline')
    t0 = time()
    cmdline_dict = dict()
    for pid in os.listdir('/proc'):
        if pid[0].isdecimal() is False or pid is '1' or pid == self_pid or pid_to_state(
                pid) != 'T':
            continue
        cmdline = pid_to_cmdline(pid)
        if cmdline == '':
            continue
        if cmdline not in cmdline_dict:
            cmdline_dict[cmdline] = 1
        else:
            cmdline_dict[cmdline] += 1
    x = cmdline_dict.items()
    y = sorted(x, key=itemgetter(1))
    print(y[::-1][:9])
    print(time() - t0)
    try:
        return y[-1][0]
    except IndexError:
        return ''


def find_bomb_uid():
    print('Find bomb eUID')
    t0 = time()
    uid_dict = dict()
    for pid in os.listdir('/proc'):
        if pid[0].isdecimal() is False or pid is '1' or pid == self_pid or pid_to_state(
                pid) != 'T':
            continue
        uid = pid_to_uid(pid)
        if uid == '':
            continue
        if uid not in uid_dict:
            uid_dict[uid] = 1
        else:
            uid_dict[uid] += 1
    x = uid_dict.items()
    y = sorted(x, key=itemgetter(1))
    print(y[::-1][:9])
    print(time() - t0)
    return y[-1][0]


def find_bomb_ppid():
    print('Find bomb PPid')
    t0 = time()
    ppid_dict = dict()
    for pid in os.listdir('/proc'):
        if pid[0].isdecimal() is False or pid is '1' or pid == self_pid or pid_to_state(
                pid) != 'T':
            continue
        ppid = pid_to_ppid(pid)
        if ppid == '':
            continue
        if ppid not in ppid_dict:
            ppid_dict[ppid] = 1
        else:
            ppid_dict[ppid] += 1
    x = ppid_dict.items()
    y = sorted(x, key=itemgetter(1))
    print('Number of unique PPids:', len(y))
    print(y[::-1][:9])
    print(time() - t0)
    return y[-1][0]


def stop_the_world1(bomb_name):
    stop_counter = 0
    for pid in os.listdir('/proc'):
        if pid[0].isdecimal() is False:
            continue
        name = pid_to_name(pid)
        if name == bomb_name:
            try:
                os.kill(int(pid), SIGSTOP)
            except FileNotFoundError:
                pass
            except ProcessLookupError:
                pass
            stop_counter += 1
    return stop_counter


def stop_the_world(bomb_name):
    print('Send SIGSTOP to bomb processes:', bomb_name)
    time0 = time()
    x0 = -1
    while True:
        x = stop_the_world1(bomb_name)
        if x == x0:
            print('Total stop time:', round(time() - time0, 3))
            return None
        x0 = x


def kill_stopped(bomb_name):
    print('Send SIGKILL to stopped processes')
    t0 = time()
    kill_counter = 0
    for pid in os.listdir('/proc'):

        if pid[0].isdecimal() is False:
            continue

        name = pid_to_name(pid)

        try:
            state = pid_to_state(pid)
        except FileNotFoundError:
            state = ''
        except ProcessLookupError:
            state = ''

        cmdline = pid_to_cmdline(pid)

        if name == bomb_name and cmdline == bomb_cmdline and state == 'T':

            try:
                os.kill(int(pid), SIGKILL)
                kill_counter += 1
            except FileNotFoundError:
                pass
            except ProcessLookupError:
                pass

    print(
        'Number of killed processes:',
        kill_counter,
        '\nKill time:',
        round(
            time() - t0,
            3))


def cont_stopped(bomb_name):
    print('Send SIGCONT to innocent victims')
    t0 = time()
    cont_counter = 0
    for pid in os.listdir('/proc'):
        if pid[0].isdecimal() is False:
            continue
        state = pid_to_state(pid)
        name = pid_to_name(pid)

        if state == 'T' and name == bomb_name:
            try:
                cmdline = pid_to_cmdline(pid)
                os.kill(int(pid), SIGCONT)
                cont_counter += 1
                print(' ', pid, name, [cmdline])
            except FileNotFoundError:
                pass
            except ProcessLookupError:
                pass

    print('Number of continued:', cont_counter)


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

pid_set = set(os.listdir('/proc'))

print('Fork bomb detecting started!')

stdout.flush()


while True:
    stdout.flush()
    sleep(TIME)
    new_set = set(os.listdir('/proc'))
    delta = new_set - pid_set
    pid_set = new_set

    if len(delta) > PRINT_NPROC:
        print('Spawned {} processes in last {} sec'.format(len(delta), TIME))


    if len(delta) > PROC_NUM:
        t0 = time()
        print(
            'Fork bomb detected! Spawned {} processess in {} seconds'.format(
                len(delta), TIME))
        bomb_name = find_bomb_pattern(delta)
        stop_the_world(bomb_name)
        print('Analyze stopped processes...')
        bomb_cmdline = find_bomb_cmdline()
        bomb_uid = find_bomb_uid()
        bomb_ppid = find_bomb_ppid()
        if bomb_ppid != '1':
            print(
                'Try to killing bomb parent: {}, {}'.format(
                    bomb_ppid,
                    pid_to_name(bomb_ppid)))
            os.kill(int(bomb_ppid), SIGKILL)
        print("Analysis complete. It's time to kill!")
        kill_stopped(bomb_name)
        cont_stopped(bomb_name)

        print(
            'Fork bomb is deactivated in {} seconds'.format(
                round(time() - t0, 3)),
            '\n  Bomb name:   ', bomb_name,
            '\n  Bomb cmdline:', bomb_cmdline,
            '\n  Bomb eUID:   ', bomb_uid,
            '\n  Bomb PPid:   ', bomb_ppid, '\n', '=' * 78)

