log_consumer.py

"""Log consumers are responsible for fetching chia logs
and propagating them to subscribers for further handling.

This abstraction should provide an easy ability to switch between
local file reader and fetching logs from a remote machine.
The latter has not been implemented yet. Feel free to add it.
"""

# std
import logging
from abc import ABC, abstractmethod
from pathlib import Path, PurePosixPath, PureWindowsPath, PurePath
from tempfile import mkdtemp
from threading import Thread
from time import sleep
from typing import List, Tuple

# project
from src.util import OS

# lib
import paramiko
import confuse
from confuse import ConfigView
from paramiko.channel import ChannelStdinFile, ChannelStderrFile, ChannelFile
from pygtail import Pygtail  # type: ignore
from retry import retry


# Define the minimum valid 'chia_logs' config sections as needed by the log consumers
file_log_consumer_template = {
    "enable": bool,
    "file_path": confuse.Path(),
}
network_log_consumer_template = {
    "enable": bool,
    "remote_file_path": confuse.Path(),
    "remote_host": str,
    "remote_user": str,
    "remote_port": int,
}


class LogConsumerSubscriber(ABC):
    """Interface for log consumer subscribers (i.e. handlers)"""

    @abstractmethod
    def consume_logs(self, logs: str):
        """This method will be called when new logs are available"""
        pass


class LogConsumer(ABC):
    """Abstract class providing common interface for log consumers"""

    def __init__(self):
        self._subscribers: List[LogConsumerSubscriber] = []

    @abstractmethod
    def stop(self):
        pass

    def subscribe(self, subscriber: LogConsumerSubscriber):
        self._subscribers.append(subscriber)

    def _notify_subscribers(self, logs: str):
        for subscriber in self._subscribers:
            subscriber.consume_logs(logs)


class FileLogConsumer(LogConsumer):
    def __init__(self, log_path: Path):
        super().__init__()
        self._expanded_log_path = str(log_path.expanduser())
        self._offset_path = mkdtemp() / Path("debug.log.offset")
        logging.debug(f"Using temporary directory {self._offset_path} for FileLogConsumer")
        self._is_running = True
        self._thread = Thread(target=self._consume_loop)
        self._thread.start()
        self._log_size = 0

    def stop(self):
        logging.info("Stopping")

        # Cleanup the temporary file
        if self._offset_path.exists():
            logging.debug(f"Deleting {self._offset_path}")
            self._offset_path.unlink()

        self._is_running = False

    @retry((FileNotFoundError, PermissionError), delay=2)
    def _consume_loop(self):
        while self._is_running:
            sleep(1)  # throttle polling for new logs
            for log_line in Pygtail(self._expanded_log_path, read_from_end=True, offset_file=self._offset_path):
                self._notify_subscribers(log_line)


class NetworkLogConsumer(LogConsumer):
    """Consume logs over SSH from a remote harvester"""

    def __init__(
        self, remote_log_path: PurePath, remote_user: str, remote_host: str, remote_port: int, remote_platform: OS
    ):
        super().__init__()

        self._remote_user = remote_user
        self._remote_host = remote_host
        self._remote_port = remote_port
        self._remote_log_path = remote_log_path
        self._remote_platform = remote_platform
        self._log_size = 0

        self._ssh_client = paramiko.client.SSHClient()
        self._ssh_client.load_system_host_keys()
        self._ssh_client.connect(hostname=self._remote_host, username=self._remote_user, port=self._remote_port)

        # Start thread
        self._is_running = True
        self._thread = Thread(target=self._consume_loop)
        self._thread.start()

    def stop(self):
        logging.info("Stopping")
        self._is_running = False

    def _consume_loop(self):
        logging.info(
            f"Consuming remote log file {self._remote_log_path}"
            + f" from {self._remote_host}:{self._remote_port} ({self._remote_platform})"
        )


class PosixNetworkLogConsumer(NetworkLogConsumer):
    """Consume logs over SSH from a remote Linux/MacOS harvester"""

    def __init__(
        self, remote_log_path: PurePath, remote_user: str, remote_host: str, remote_port: int, remote_platform: OS
    ):
        logging.info("Enabled Posix network log consumer.")
        super(PosixNetworkLogConsumer, self).__init__(
            remote_log_path, remote_user, remote_host, remote_port, remote_platform
        )

    def _consume_loop(self):
        super(PosixNetworkLogConsumer, self)._consume_loop()

        stdin, stdout, stderr = self._ssh_client.exec_command(f"tail -F {self._remote_log_path}")

        while self._is_running:
            log_line = stdout.readline()
            self._notify_subscribers(log_line)


class WindowsNetworkLogConsumer(NetworkLogConsumer):
    """Consume logs over SSH from a remote Windows harvester"""

    def __init__(
        self, remote_log_path: PurePath, remote_user: str, remote_host: str, remote_port: int, remote_platform: OS
    ):
        logging.info("Enabled Windows network log consumer.")
        super(WindowsNetworkLogConsumer, self).__init__(
            remote_log_path, remote_user, remote_host, remote_port, remote_platform
        )

    def _consume_loop(self):
        super(WindowsNetworkLogConsumer, self)._consume_loop()

        stdin, stdout, stderr = self._read_log()

        while self._is_running:
            if self._has_rotated(self._remote_log_path):
                sleep(1)
                stdin, stdout, stderr = self._read_log()

            log_line = stdout.readline()
            self._notify_subscribers(log_line)

    def _read_log(self) -> Tuple[ChannelStdinFile, ChannelFile, ChannelStderrFile]:
        stdin, stdout, stderr = self._ssh_client.exec_command(
            f"powershell.exe Get-Content {self._remote_log_path} -Wait -Tail 1"
        )

        return stdin, stdout, stderr

    def _has_rotated(self, path: PurePath) -> bool:
        stdin, stdout, stderr = self._ssh_client.exec_command(f"powershell.exe Write-Host(Get-Item {str(path)}).length")

        old_size = self._log_size
        self._log_size = int(stdout.readline())

        return old_size > self._log_size


def get_host_info(host: str, user: str, path: str, port: int) -> Tuple[OS, PurePath]:
    client = paramiko.client.SSHClient()
    client.load_system_host_keys()
    client.connect(hostname=host, username=user, port=port)

    stdin, stdout, stderr = client.exec_command("uname -a")
    fout: str = stdout.readline().lower()
    ferr: str = stderr.readline().lower()

    if "linux" in fout:
        return OS.LINUX, PurePosixPath(path)
    elif "darwin" in fout:
        return OS.MACOS, PurePosixPath(path)
    elif "not recognized" in ferr:
        return OS.WINDOWS, PureWindowsPath(path)
    else:
        logging.error("Found unsupported platform on remote host, assuming Linux and hope for the best.")

    return OS.LINUX, PurePosixPath(path)


def create_log_consumer_from_config(config: ConfigView) -> LogConsumer:
    enabled_consumer = None
    for consumer in config.keys():
        if config[consumer]["enable"].get(bool):
            if enabled_consumer:
                logging.error("Detected multiple enabled consumers. This is unsupported configuration!")
                return None
            enabled_consumer = consumer
    if enabled_consumer is None:
        logging.critical("Couldn't find enabled log consumer in config.yaml")
        exit(1)

    enabled_consumer_config = config[enabled_consumer]

    if enabled_consumer == "file_log_consumer":
        # Validate config against template
        valid_config = enabled_consumer_config.get(file_log_consumer_template)
        log_path = valid_config["file_path"]
        logging.info(f"Consuming logs locally from {log_path}")
        return FileLogConsumer(log_path=log_path)

    if enabled_consumer == "network_log_consumer":
        # Validate config against template
        valid_config = enabled_consumer_config.get(network_log_consumer_template)
        remote_port = valid_config["remote_port"]
        remote_user = valid_config["remote_user"]
        remote_host = valid_config["remote_host"]
        remote_path = valid_config["remote_file_path"]

        logging.info(f"Consuming logs remotely from {remote_user}@{remote_host}:{remote_port}:{remote_path}")

        platform, path = get_host_info(remote_host, remote_user, remote_path, remote_port)

        if platform == OS.WINDOWS:
            return WindowsNetworkLogConsumer(
                remote_log_path=path,
                remote_host=valid_config["remote_host"],
                remote_user=valid_config["remote_user"],
                remote_port=remote_port,
                remote_platform=platform,
            )
        else:
            return PosixNetworkLogConsumer(
                remote_log_path=path,
                remote_host=valid_config["remote_host"],
                remote_user=valid_config["remote_user"],
                remote_port=remote_port,
                remote_platform=platform,
            )

    logging.critical("Unknown log consumer type enabled, typo?")
    exit(1)