#!/opt/alt/python35/bin/python3.5

import json
import logging
import logging.handlers
import os.path
import socket
import subprocess
import sys
import time
from typing import Optional, List

import sentry_sdk
from sentry_sdk import configure_scope, push_scope

logging.raiseExceptions = False

CONNECT_TIMEOUT = 10
LICENSE = '/var/imunify360/license.json'
REQUEST_TIMEOUT = 60
RETRY_DELAY = 10
IMUNIFY360 = 'imunify360'
SOCKET_PATH = '/var/run/defence360agent/simple_rpc.sock'
SYSTEMCTL = '/usr/bin/systemctl'
SERVICE = 'service'
RESTART = 'restart'
STATUS = 'status'
SENTRY_DSN = 'https://831a73da71fc4a6abe776f7d0adcc48d:3b9c6c32229744388f98e85dc9104da2@sentry.cloudlinux.com/3?timeout=20'  # noqa: E501
IMUNIFY360_PKG = 'imunify360-firewall'


def service_is_running(has_systemctl: bool, name: str) -> bool:
    if has_systemctl:
        cmd = [SYSTEMCTL, STATUS, name]
    else:
        cmd = [SERVICE, name, STATUS]
    cp = subprocess.run(cmd, stdin=subprocess.DEVNULL,
                        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                        check=False)
    return cp.returncode == 0


def restart_service(has_systemctl: bool, name: str) -> None:
    if has_systemctl:
        cmd = [SYSTEMCTL, RESTART, name]
    else:
        cmd = [SERVICE, name, RESTART]
    subprocess.run(cmd, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL,
                   stderr=subprocess.DEVNULL, check=False)


def get_server_id() -> str:
    try:
        return json.load(open(LICENSE))['id']
    except Exception:
        return 'UNKNOWN'


def _init_sentry():
    sentry_sdk.init(dsn=SENTRY_DSN)
    with configure_scope() as scope:
        scope.user = {'id': get_server_id()}


def setup_logging() -> logging.Logger:
    logger = logging.getLogger(os.path.basename(sys.argv[0]))
    logger.setLevel(logging.INFO)
    handler = logging.handlers.SysLogHandler('/dev/log')
    formatter = logging.Formatter('%(name)s: %(message)s')
    handler.formatter = formatter
    logger.addHandler(handler)
    _init_sentry()
    return logger


def rpc_request(*args, **kwargs):
    r = send_to_agent_socket(list(args), kwargs)
    if r['result'] != 'success':
        raise ValueError(r.get('messages', 'Unknown error'))
    return r.get('data')


def send_to_agent_socket(command: list, params: dict):
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
        sock.settimeout(CONNECT_TIMEOUT)
        sock.connect(SOCKET_PATH)
        msg = json.dumps({'command': command, 'params': params}) + '\n'
        start_time = time.monotonic()
        sock.settimeout(REQUEST_TIMEOUT)
        sock.sendall(msg.encode())

        remaining_time = start_time + REQUEST_TIMEOUT - time.monotonic()
        if remaining_time <= 0:
            raise socket.timeout()
        sock.settimeout(remaining_time)
        with sock.makefile(encoding='utf-8') as file:
            response = file.readline()
        if not response:
            raise ValueError('Empty response from socket')
        return json.loads(response)


def rpc_request_with_retries(rpc_timeout: int, logger) -> Optional[dict]:
    start = time.time()
    while time.time() - start < rpc_timeout:
        try:
            return rpc_request('health')
        except Exception:
            logger.exception('Failed to perform RPC request')
            time.sleep(RETRY_DELAY)
    return None


def systemctl_present() -> bool:
    return os.path.isfile(SYSTEMCTL)


def collect_output(cmd: List[str]) -> str:
    try:
        cp = subprocess.run(cmd, stdin=subprocess.DEVNULL,
                            stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
    except OSError:
        return ''
    if cp.returncode != 0:
        return ''
    return cp.stdout.decode()


def get_rpm_version(pkg: str) -> str:
    cmd = ['rpm', '-q', '--queryformat=%{VERSION}-%{RELEASE}', pkg]
    return collect_output(cmd)


def get_dpkg_version(pkg: str) -> str:
    cmd = ['xdpkg-query', '--showformat=${Version}', '--show', pkg]
    return collect_output(cmd)


def imunify360_version() -> str:
    version = get_rpm_version(IMUNIFY360_PKG)
    if not version:
        version = get_dpkg_version(IMUNIFY360)
    return version


def main(rpc_timeout):
    logger = setup_logging()
    has_systemctl = systemctl_present()
    if not service_is_running(has_systemctl, IMUNIFY360):
        logger.info('%s is not running', IMUNIFY360)
        return
    response = rpc_request_with_retries(rpc_timeout, logger)
    if response is None:
        with push_scope() as scope:
            scope.set_tag('version', imunify360_version())
            logger.error('Restarting due to RPC failures')
        restart_service(has_systemctl, IMUNIFY360)
        return
    with configure_scope() as scope:
        scope.set_tag('version', response.get('version', 'UNKNOWN'))
    if not response.get('healthy', False):
        logger.error('Restarting due to health report: %s',
                     response.get('why'))
        restart_service(has_systemctl, IMUNIFY360)
    else:
        logger.info('%s is healthy: %s', IMUNIFY360, response.get('why'))


if __name__ == '__main__':
    rpc_timeout = int(sys.argv[1])
    main(rpc_timeout)
