#
# Copyright (C) 2018 Andreas Gerstmayr <andreas@gerstmayr.me>
# Based on the execsnoop BCC tool by Brendan Gregg:
# https://github.com/iovisor/bcc/blob/master/tools/execsnoop.py
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
""" PCP BCC PMDA execsnoop module """

# pylint: disable=invalid-name, too-few-public-methods, too-many-instance-attributes

import ctypes as ct
from collections import deque, defaultdict
from threading import Lock, Thread
from os import path
import sys
import re

from bcc import BPF

from pcp.pmapi import pmUnits
from cpmapi import PM_TYPE_32, PM_TYPE_U32, PM_TYPE_STRING, PM_SEM_INSTANT
from cpmda import PMDA_FETCH_NOVALUES

from modules.pcpbcc import PCPBCCBase

#
# BPF program
#
bpf_src = "modules/execsnoop.bpf"

#
# PCP BCC PMDA constants
#
MODULE = 'execsnoop'
BASENS = 'proc.exec.'
units_none = pmUnits(0, 0, 0, 0, 0, 0)

class EventType(object):
    """ Event type """
    EVENT_ARG = 0
    EVENT_RET = 1

#
# PCP BCC Module
#
class PCPBCCModule(PCPBCCBase):
    """ PCP BCC execsnoop module """
    def __init__(self, config, log, err, _):
        """ Constructor """
        PCPBCCBase.__init__(self, MODULE, config, log, err)

        self.include_failed = False
        self.command = None
        self.args = None
        self.max_args = 20
        self.process_count = 20
        self.buffer_page_count = 64

        for opt in self.config.options(MODULE):
            if opt == 'include_failed':
                self.include_failed = self.config.getboolean(MODULE, opt)
            if opt == 'command':
                self.command = self.config.get(MODULE, opt).encode(sys.getfilesystemencoding())
            if opt == 'args':
                self.args = self.config.get(MODULE, opt).encode(sys.getfilesystemencoding())
            if opt == 'max_args':
                self.max_args = int(self.config.get(MODULE, opt))
            if opt == 'process_count':
                self.process_count = int(self.config.get(MODULE, opt))
            if opt == 'buffer_page_count':
                self.buffer_page_count = int(self.config.get(MODULE, opt))
                if not self.buffer_page_count or \
                   self.buffer_page_count & (self.buffer_page_count - 1):
                    raise RuntimeError("Buffer page count is not power of two.")

        self.cache = deque(maxlen=self.process_count)
        self.insts = {str(i) : ct.c_int(1) for i in range(self.process_count)}

        self.lock = Lock()
        self.argv_cache = defaultdict(list)
        self.thread = Thread(name="bpfpoller", target=self.perf_buffer_poller)
        self.thread.daemon = True

        # Exit hard if impossible to continue
        if self.bcc_version() == "0.5.0":
            raise RuntimeError("BCC 0.5.0 is too old for this module to work properly.")

        self.log("Initialized.")

    def metrics(self):
        """ Get metric definitions """
        name = BASENS
        self.items = (
            # Name - reserved - type - semantics - units - help
            (name + 'comm', None, PM_TYPE_STRING, PM_SEM_INSTANT, units_none, 'command'),
            (name + 'pid', None, PM_TYPE_U32, PM_SEM_INSTANT, units_none, 'PID'),
            (name + 'ppid', None, PM_TYPE_U32, PM_SEM_INSTANT, units_none, 'parent PID'),
            (name + 'ret', None, PM_TYPE_32, PM_SEM_INSTANT, units_none, 'return code'),
            (name + 'args', None, PM_TYPE_STRING, PM_SEM_INSTANT, units_none, 'command arguments'),
        )
        return True, self.items

    @staticmethod
    def get_ppid(pid):
        """ Get parent PID """
        try:
            with open("/proc/%d/status" % pid) as status:
                for line in status:
                    if line.startswith("PPid:"):
                        return int(line.split()[1])
        except IOError:
            pass
        return 0

    def handle_event(self, _cpu, data, _size):
        """ Event handler """
        event = self.bpf["events"].event(data)
        skip = False

        if event.type == EventType.EVENT_ARG:
            self.argv_cache[event.pid].append(event.argv)
        elif event.type == EventType.EVENT_RET:
            if event.retval != 0 and not self.include_failed:
                skip = True
            if self.command and not re.search(bytes(self.command), event.comm):
                skip = True
            if self.args and not re.search(bytes(self.args), b" ".join(self.argv_cache[event.pid])):
                skip = True

            if not skip:
                ppid = event.ppid if event.ppid > 0 else self.get_ppid(event.pid)
                argv_text = b" ".join(self.argv_cache[event.pid]).replace(b"\n", b"\\n")

                self.lock.acquire()
                self.cache.appendleft([
                    event.comm.decode(), event.pid, ppid, event.retval, argv_text.decode()
                ])
                self.lock.release()
            try:
                del self.argv_cache[event.pid]
            except Exception: # pylint: disable=broad-except
                pass

    def perf_buffer_lost_cb(self, lost_cnt):
        """ Callback for lost perf buffer events """
        self.err("Lost %d events; buffer_page_count should be increased." % lost_cnt)

    def compile(self):
        """ Compile BPF """
        try:
            if not self.bpf_text:
                with open(path.dirname(__file__) + '/../' + bpf_src) as src:
                    self.bpf_text = src.read()

                self.bpf_text = self.bpf_text.replace("MAXARG", str(self.max_args))

            bpf_text = self.bpf_text
            bpf_text = bpf_text.replace('UID_FILTER', '')
            bpf_text = bpf_text.replace('container_should_be_filtered()', '0')

            if self.debug:
                self.log("BPF to be compiled:\n" + bpf_text.strip())

            self.bpf = BPF(text=bpf_text)
            execve_fnname = self.get_syscall_fnname("execve")
            self.bpf.attach_kprobe(event=execve_fnname, fn_name="syscall__execve")
            self.bpf.attach_kretprobe(event=execve_fnname, fn_name="do_ret_sys_execve")
            self.bpf["events"].open_perf_buffer(self.handle_event,
                                                page_cnt=self.buffer_page_count,
                                                lost_cb=self.perf_buffer_lost_cb)
            self.thread.start()
            self.log("Compiled.")
        except Exception as error: # pylint: disable=broad-except
            self.bpf = None
            self.err(str(error))
            self.err("Module NOT active!")
            raise

    def refresh(self):
        """ Refresh BPF data """
        if self.bpf is None:
            return None

        return self.insts

    def bpfdata(self, item, inst):
        """ Return BPF data as PCP metric value """
        try:
            key = int(self.pmdaIndom.inst_name_lookup(inst))
            self.lock.acquire()
            value = self.cache[key][item]
            self.lock.release()
            return [value, 1]
        except Exception: # pylint: disable=broad-except
            self.lock.release()
            return [PMDA_FETCH_NOVALUES, 0]