#!/usr/bin/python3

# Copyright (c) 2009 Peter Palfrader
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


import sqlite3
import os
import os.path
import optparse
import sys
import yaml
import threading
import time
import subprocess
import time
from dbhelper import DBHelper
import smtplib
from email.mime.text import MIMEText

usage="""Usage: %prog -d <db> [options] <action>

Action is one of:
  - init        Initialize the database
  - check       run check against list of hosts
  - list-bad    List servers not to be used (down, in shutdown, ...)
  - update-bad  Update auto-dns's status directory about bad/ok servers
"""

check_timeout = 35
check_command = {
    'ping-check':     '/usr/lib/nagios/plugins/check_ping -H @@HOST@@ -w 800,40% -c 1500,60% -p 10',
    'shutdown-check': '! /usr/lib/nagios/plugins/check_nrpe -H @@HOST@@ -n -c dsa2_shutdown | grep system-in-shutdown',
    'http-check':     '/usr/lib/nagios/plugins/check_http -H @@HOST@@ -t 30 -w 15',
    'debianhealth-check': '/usr/lib/nagios/plugins/check_http -I @@HOST@@ -u http://debian.backend.mirrors.debian.org/_health -t 30 -w 15',
    'debughealth-check': '/usr/lib/nagios/plugins/check_http -I @@HOST@@ -u http://debug.backend.mirrors.debian.org/_health -t 30 -w 15',
}
expire_results_time = 900
mail_from = 'tor-misc@commit.noreply.org'
mail_rcpt = 'tor-misc@commit.noreply.org'
mail_nsa_project = 'tor-nagios'


def init(options):
    if os.path.exists(options.db):
        sys.stderr.write("File %s already exists.\n"%options.db)
        sys.exit(1)

    db = DBHelper(options.db)
    db.execute('CREATE TABLE config (key TEXT UNIQUE, value TEXT)')
    db.execute("""CREATE TABLE host_status (
                    host TEXT NOT NULL,
                    test TEXT NOT NULL,
                    ts INTEGER NOT NULL,
                    soft BOOLEAN NOT NULL,
                    hard BOOLEAN NOT NULL,
                    msg TEXT)
               """)
    db.execute('INSERT INTO config (key, value) VALUES (?,?)', ('db_revision', '1'))
    db.commit()
    db.close()


def check_one(*args, **kw):
    cmd = check_command[kw['check']].replace('@@HOST@@', kw['host'])
    p = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True, encoding='utf-8')
    (out, none) = p.communicate('')

    kw['result']['ret'] = p.returncode
    kw['result']['msg'] = out.strip("\n")

def run_checks(db, checks, hosts, check, hard = False):
    for host in hosts:
        c = {}
        c['host'] = host
        c['check'] = check
        c['hard'] = hard

        c['result'] = {}
        c['thread'] = threading.Thread(target=check_one, kwargs={ 'check': check, 'host': hosts[host][4], 'result': c['result'] })
        c['thread'].start()

        checks.append(c)

def join_checks(checks):
    start = time.time()
    for c in checks:
        time_left = max(check_timeout - (time.time() - start), 0)
        c['thread'].join(time_left)
        if c['thread'].is_alive():
            # the thread didn't terminate, the join timed out

            # make sure the result hash is re-created so that
            # thread cannot write to it after it failed to
            # deliver on time
            c['result'] = {
              'ret': 2,
              'msg': 'check timeout' }

def insert_results(db, checks):
    for c in checks:
        failed = c['result']['ret'] != 0
        db.execute('INSERT INTO host_status (host, test, ts, soft, hard, msg) VALUES (?,?,?,?,?,?)',
            ( c['host'],
              c['check'],
              int(time.time()),
              failed,
              c['hard'] and failed,
              c['result']['msg']
            ) )

def dependency_checks(hosts, db):
    bad_hosts = None
    checks = []
    for hostname, host in hosts.items():
        if not 'depends' in host: continue
        if bad_hosts is None: bad_hosts = set(get_bad_from_db(db))
        depends = set(host['depends'])

        intersection = sorted(list(bad_hosts & depends))

        c = {}
        c['host'] = hostname
        c['check'] = 'depends'
        c['hard'] = True

        c['result'] = {}
        if len(intersection) > 0:
            c['result']['ret'] = 1
            c['result']['msg'] = "Failed: %s"%(", ".join(intersection),)
        else:
            c['result']['ret'] = 0
            c['result']['msg'] = "no dependencies failed."

        checks.append(c)
    insert_results(db, checks)

def check(options):
    if not os.path.exists(options.db):
        sys.stderr.write("DB %s does not exist.\n"%options.db)
        sys.exit(1)

    if not options.hosts:
        sys.stderr.write("Need hosts.yaml file\n")
        sys.exit(1)

    hosts = yaml.safe_load(open(options.hosts).read())['hosts']
    db = DBHelper(options.db)
    checks = []
    for x in ['ping-check', 'http-check']:
        run_checks(db, checks, dict(filter(lambda h: 'checks' in h[1] and x in h[1]['checks'], hosts.items())), x)
    for x in ['shutdown-check', 'debianhealth-check', 'debughealth-check']:
        run_checks(db, checks, dict(filter(lambda h: 'checks' in h[1] and x in h[1]['checks'], hosts.items())), x, True)
    join_checks(checks)
    insert_results(db, checks)
    dependency_checks(hosts, db)
    db.commit()
    db.close()

def cleanup_bad_in_db(db):
    # delete all entries older than 15 minutes.
    db.execute("DELETE FROM host_status WHERE ts < ?", (int(time.time()) - expire_results_time, ))
    db.commit()

def get_bad_from_db(db):
    entries = db.query("""
        SELECT total, soft*1.0/total as soft, hard, host, test
          FROM (SELECT count(*) AS total, sum(soft) AS soft, sum(hard) AS hard, host, test
                  FROM host_status
                  GROUP BY host, test)
        WHERE soft*1.0/total > 0.40
          OR hard > 0
        """)

    bad_hosts = {}
    for e in entries:
        if not e['host'] in bad_hosts:
            bad_hosts[e['host']] = []

        if e['hard']:
            bad_hosts[e['host']].append(e['test'])
        elif e['soft'] > 0:
            bad_hosts[e['host']].append(e['test'] + " (%.2f%%)"%(e['soft'] * 100))

    return bad_hosts

def get_bad(options):
    if not os.path.exists(options.db):
        sys.stderr.write("DB %s does not exist.\n"%options.db)
        sys.exit(1)

    db = DBHelper(options.db)
    cleanup_bad_in_db(db)
    bad_hosts = get_bad_from_db(db)
    db.close()
    return bad_hosts


def list_bad(options):
    bad_hosts = get_bad(options)

    for host in bad_hosts:
        print("%s failed: %s"% ( host, '; '.join(bad_hosts[host])))

def notify(s, options):
    if not options.mail_only:
        print(s)

    if options.mail:
        msg = MIMEText("[mini-nag/auto-dns] " + s)
        msg['Subject'] = 'Announce %s'%mail_nsa_project
        msg['From'] = mail_from
        msg['To'] = mail_rcpt
        msg['Precedence'] = 'junk'

        s = smtplib.SMTP('127.0.0.1')
        s.sendmail(mail_from, [mail_rcpt], msg.as_string())
        s.quit()

def update_bad(options):
    bad_hosts = get_bad(options)

    if not options.dir:
        sys.stderr.write("Need status directory\n")
        sys.exit(1)

    if not os.path.isdir(options.dir):
        sys.stderr.write("%s does not exist or is not a directory\n"%(options.dir))
        sys.exit(1)

    already_bad = {}
    for d in os.listdir(options.dir):
        # . and .. do not show up in listdir() but other dotfiles may (.dummy, etc)
        if d.startswith('.'):
            continue

        if not d in bad_hosts:
            os.unlink(os.path.join(options.dir, d))
            cur_time = time.strftime("%H:%M:%S +0000", time.gmtime())
            notify("%s is considered OK again @ %s."%(d,cur_time), options)
        else:
            already_bad[d] = True

    for host in bad_hosts:
        if not host in already_bad:
            open(os.path.join(options.dir, host), 'w').close()
            cur_time = time.strftime("%H:%M:%S +0000", time.gmtime())
            notify("%s is considered BAD (%s) @ %s."%(host, '; '.join(bad_hosts[host]),cur_time), options)


parser = optparse.OptionParser(usage=usage)
parser.add_option("-d", "--db", help="database file", metavar="FILE")
parser.add_option("-H", "--hosts", help="host definition yaml file", metavar="FILE")
parser.add_option("-D", "--dir", help="status directory", metavar="DIR")
parser.add_option("-m", "--mail", help="send notifies per mail (in addition to print to stdout) on status-update", action="store_true")
parser.add_option("-M", "--mail-only", help="send notifies only per mail", action="store_true")
(options, args) = parser.parse_args()

if len(args) != 1:
    parser.print_help(file = sys.stderr)
    sys.exit(1)
if not options.db:
    parser.print_help(file = sys.stderr)
    sys.exit(1)


if args[0] == "init":
    init(options)
elif args[0] == "check":
    check(options)
elif args[0] == "list-bad":
    list_bad(options)
elif args[0] == "update-bad":
    update_bad(options)
else:
    parser.print_help(file = sys.stderr)
    sys.exit(1)



#if not os.path.exists(db):
#    conn = sqlite3.connect('status.db')

# vim: ts=4 sw=4 et: