#! /usr/bin/python3

import os
import requests
import time
import calendar
import logging
import subprocess
from email.utils import parsedate
import urllib3.util.connection

logging.basicConfig(level=logging.INFO)

if 'MIRROR_CHECK_HOSTS' in os.environ:
    HOSTS = os.environ['MIRROR_CHECK_HOSTS'].split()
else:
    HOSTS = open(os.environ['MIRROR_CHECK_HOSTS_FILE']).read().split()

OUTPUT_DIR = "/run/dsa-mirror-health-{}".format(os.environ['MIRROR_CHECK_SERVICE'])
HEALTH_FILE = os.path.join(OUTPUT_DIR, "health")
URL = os.environ['MIRROR_CHECK_URL']
HEALTH_CHECK_URL = os.environ['MIRROR_CHECK_HEALTH_URL']
INTERVAL = int(os.environ.get('MIRROR_CHECK_INTERVAL', '60'))

def retrieve_from_host(host, url):
    orig_create_connection = urllib3.util.connection.create_connection
    def patched_create_connection(address, *args, **kwargs):
        _host, port = address
        return orig_create_connection((host, port), *args, **kwargs)
    headers = {'User-Agent': 'mirror-health'}
    urllib3.util.connection.create_connection = patched_create_connection
    try:
        return requests.get(url, headers=headers, timeout=5, allow_redirects=False)
    finally:
        urllib3.util.connection.create_connection = orig_create_connection

def last_modified(response):
    lm = 0
    if response.status_code == 200 and response.headers.get('last-modified'):
        lm = calendar.timegm(parsedate(response.headers['last-modified']))
    return lm

def healthy(response):
    if response.status_code == 200:
        return True
    return False

def check_shutdown():
    if os.path.exists('/run/systemd/shutdown/scheduled'):
        logging.info("considering myself unhealthy, shutdown scheduled")
        return False
    return True

def check_uptodate():
    latest_ts = 0
    for host in HOSTS:
        try:
            lm = last_modified(retrieve_from_host(host, URL))
            logging.debug("lm for host %s: %s", host, lm)
            if healthy(retrieve_from_host(host, HEALTH_CHECK_URL)):
                latest_ts = max(latest_ts, lm)
        except (requests.exceptions.ProxyError, requests.exceptions.ReadTimeout, requests.exceptions.ConnectTimeout, requests.exceptions.ConnectionError):
            pass
    try:
        local_lm = last_modified(retrieve_from_host('localhost', URL))
    except (requests.exceptions.ProxyError, requests.exceptions.ReadTimeout, requests.exceptions.ConnectTimeout, requests.exceptions.ConnectionError):
        return False
    logging.debug("lm for localhost: %s", local_lm)
    if local_lm < latest_ts:
        logging.info("considering myself unhealthy my ts=%s latest_ts=%s", local_lm, latest_ts)
        return False
    return True

while True:
    start = time.time()
    if check_shutdown() and check_uptodate():
        logging.info("considering myself healthy")
        open(HEALTH_FILE, 'w').write("OK")
    else:
        try:
            os.remove(HEALTH_FILE)
        except OSError:
            pass
    sleep_time = start + INTERVAL - time.time()
    logging.debug("sleeping for %d seconds", sleep_time)
    time.sleep(sleep_time)