#!/usr/bin/env python
#
# This script is here because 800.loginfail does not check for log record year,
# and auth.log can miss rotation and therefore include records from last year.
# (see https://bugs.freenas.org/issues/18615)

import bz2
from datetime import datetime, timedelta
import glob
import gzip
import os
import re
import sys


def catmsgs():
    log_directory = "/var/log"

    for log_file in sorted(
        filter(
            lambda path: re.match(".*\.[0-9]+\.[^.]+$", path),
            glob.glob("/var/log/auth.log.*.*"),
        ),
        key=lambda path: int(path.split(".")[-2]),
        reverse=True,
    ):
        if log_file.endswith(".bz2"):
            try:
                with bz2.BZ2File(log_file, "rb") as f:
                    yield from f
            except IOError:
                pass

        if log_file.endswith(".gz"):
            try:
                with gzip.GzipFile(log_file, "rb") as f:
                    yield from f
            except IOError:
                pass

    try:
        with open(os.path.join(log_directory, "auth.log"), "rb") as f:
            yield from f
    except IOError:
        pass


def get_login_failures(now, messages):
    """
    >>> get_login_failures(datetime(year=2017, month=8, day=31), [
    ...    b'Aug 30 invalid login\\n',  # 2017
    ...    b'Aug 31 invalid login\\n',  # 2017
    ... ])
    [b'Aug 30 invalid login\\n']

    >>> get_login_failures(datetime(year=2017, month=8, day=31), [
    ...    b'Aug 30 invalid login\\n',  # 2016
    ...    b'Oct 18 invalid login\\n',  # 2016
    ...    b'Aug 31 invalid login\\n',  # 2017
    ... ])
    []

    >>> get_login_failures(datetime(year=2017, month=8, day=31), [
    ...    b'Aug 30 invalid login\\n',  # 2015
    ...    b'Oct 18 invalid login\\n',  # 2015
    ...    b'Aug 31 invalid login\\n',  # 2016
    ...    b'Aug 30 bad login\\n',      # 2017
    ... ])
    [b'Aug 30 bad login\\n']

    >>> get_login_failures(datetime(year=2017, month=8, day=31), [
    ...    b'Aug 30 invalid login\\n',  # 2017
    ...    b'Aug 31 invalid login\\n',  # 2017
    ...    b'\\n',                      # Random empty line at the end of file
    ... ])
    [b'Aug 30 invalid login\\n']
    """

    yesterday = (now - timedelta(days=1)).strftime("%b %e ").encode("ascii")
    today = now.strftime("%b %e ").encode("ascii")

    login_failures = []
    for message in messages:
        if message.strip():
            if message.startswith(yesterday):
                if re.search(rb"\b(fail(ures?|ed)?|invalid|bad|illegal|auth.*error)\b", message):
                    login_failures.append(message)

            if not message.startswith(yesterday) and not message.startswith(today):
                login_failures = []

    return login_failures

if __name__ == "__main__":
    # import doctest
    # doctest.testmod()
    # sys.exit(0)

    # import pprint
    # pprint.pprint(list(catmsgs()))
    # sys.exit(0)

    login_failures = get_login_failures(datetime.now(), catmsgs())
    os.write(1, b"\n%s login failures:\n" % os.environ.get("host", "").encode("utf-8"))
    os.write(2, b"".join(login_failures))
    if login_failures:
        sys.exit(1)
    else:
        sys.exit(0)
