#!/usr/bin/env python3

import requests
import argparse
import json
import ssl
import hashlib
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from urllib3.poolmanager import PoolManager

def fingerprint_checking_SSLSocket(_fingerprint:str):
    class SSLSocket(ssl.SSLSocket):
        fingerprint = _fingerprint.replace(":", "").lower()

        def do_handshake(self, *args, **kw):
            res = super().do_handshake(*args, **kw)

            # Get full certificate in DER format
            der_bytes = self.getpeercert(binary_form=True)

            crt_sha256 = hashlib.sha256(der_bytes).hexdigest()

            if crt_sha256.lower() != self.fingerprint.lower():
                raise ssl.SSLError(
                    "Server %r certificate fingerprint (sha1) %s does not match %r"
                        % (
                            self.server_hostname,
                            crt_sha256,
                            self.fingerprint,
                        )
                )

            return res

    return SSLSocket

def api_get(session, url):
    try:
        r = session.get(url, timeout=50)

        if r.status_code == 404:
            raise RuntimeError(f"API endpoint not found (404): {url}")

        if not r.ok:
            raise RuntimeError(f"API error {r.status_code}: {url} - {r.text}")

        return r.json().get("data", {})

    except ConnectionError as e:
        raise RuntimeError(
            f"Connection refused to {url} (host or port unreachable)"
        ) from e

    except requests.exceptions.Timeout as e:
        raise RuntimeError(
            f"Timeout while connecting to {url}"
        ) from e

    except requests.exceptions.HTTPError as e:
        raise RuntimeError(
            f"HTTP error from {url}: {e}"
        ) from e

    except requests.exceptions.RequestException as e:
        raise RuntimeError(
            f"Request failed for {url}: {e}"
        ) from e
    
def create_session():
    session = requests.Session()

    retry = Retry(
        total=3,                 # total retries
        connect=3,               # connection retries
        read=3,                  # read retries
        backoff_factor=0.5,      # 0.5s, 1s, 2s
        status_forcelist=[500, 502, 503, 504],
        allowed_methods=["GET"],
        raise_on_status=False,
    )

    adapter = HTTPAdapter(
        max_retries=retry,
        pool_connections=5,
        pool_maxsize=5,
    )

    session.mount("https://", adapter)
    session.mount("http://", adapter)

    return session

class FingerprintAdapter(HTTPAdapter):
    def init_poolmanager(self, connections, maxsize, block=False, **kwargs):
        ctx = ssl.create_default_context()
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE
        ctx.sslsocket_class = fingerprint_checking_SSLSocket(PBS_FINGERPRINT)

        self.poolmanager = PoolManager(
            num_pools=connections,
            maxsize=maxsize,
            block=block,
            ssl_context=ctx,
            assert_hostname=False,
        )

parser = argparse.ArgumentParser("agent_pbs_special_agent")
parser.add_argument(
    "--host",
    help="PBS Host address or fqdn without https:// and port.",
    type=str,
    default="pbs")
parser.add_argument(
    "--port",
    help="PBS https port.",
    type=int,
    default="8007")
parser.add_argument(
    "--fingerprint",
    help="Fingerprint of the PBS if not using valid cert.",
    type=str,
    default="")
parser.add_argument(
    "--tokenid",
    help="API Token with audit permissions.",
    type=str,
    default="admin@pbs!checkmk")
parser.add_argument(
    "--secret",
    help="Secret for API Token.",
    type=str)
args = parser.parse_args()

PBS_HOST = args.host
API_TOKEN_ID = args.tokenid
PBS_PORT = args.port
PBS_FINGERPRINT = args.fingerprint
API_TOKEN_SECRET = args.secret

headers = {
    "Authorization": f"PBSAPIToken={API_TOKEN_ID}:{API_TOKEN_SECRET}"
}
# Disable SSL verification only if you use self-signed certs
session = create_session()
if PBS_FINGERPRINT != "":
    session.mount("https://", FingerprintAdapter())
# session.verify = False
session.headers.update(headers)

API_URL = f"https://{PBS_HOST}:{PBS_PORT}"
try:
    return_json = {}
    return_json["version"] = api_get(session, f"{API_URL}/api2/json/version")
    return_json["status"] = api_get(session, f"{API_URL}/api2/json/nodes/localhost/status")
    return_json["tasks"] = api_get(session, f"{API_URL}/api2/json/nodes/localhost/tasks")
    return_json["ds_usage"] = api_get(session, f"{API_URL}/api2/json/status/datastore-usage")
    return_json["sync"] = api_get(session, f"{API_URL}/api2/json/admin/sync")
    return_json["gc"] = api_get(session, f"{API_URL}/api2/json/admin/gc")
    return_json["prune"] = api_get(session, f"{API_URL}/api2/json/admin/prune")
    return_json["verify"] = api_get(session, f"{API_URL}/api2/json/admin/verify")
    for d in return_json["ds_usage"]:
        # we do not need history data
        d.pop("history", None)
    print('<<<pbs_special_agent:sep(0)>>>')
    print(json.dumps(return_json))
    exit(0)
except Exception as e:
    print('<<<pbs_special_agent:sep(0)>>>')
    print(json.dumps({"error": str(e)}))
    exit(0)


