Checks domains against ssllabs.

ssllabs_check.py

#!/usr/bin/env python3
"""
SSL Labs Testing API

https://github.com/ssllabs/ssllabs-scan/blob/master/ssllabs-api-docs-v3.md
"""
import logging
import re
import sys
import time
from typing import List, Optional

import attr
import click
import jinja2
import requests
import jmespath

logging.basicConfig(format="[%(levelname)-8s] %(message)s", level=logging.INFO)
log = logging.getLogger(__name__)

_templates = {
    "report": """\
{% for check in manager.checks -%}
{{ check.host }}: {{ check.result | tojson(indent=2) }}

{% endfor -%}
"""
}


@attr.s(slots=True)
class SslLabsCheck:
    api_base_url = "https://api.ssllabs.com/api/v3/"
    assert api_base_url.endswith("/")
    status_values = ("DNS", "ERROR", "IN_PROGRESS", "READY")

    host = attr.ib()
    fresh = attr.ib(default=False)
    status = attr.ib(init=False, default=None)
    result = attr.ib(init=False, default=None)

    @staticmethod
    def _call(verb, params=None):
        resp = requests.get(SslLabsCheck.api_base_url + verb, params=params)
        if resp.status_code == 429:
            # back off a little
            time.sleep(5)
            return SslLabsCheck._call(verb, params)
        resp.raise_for_status()
        result = resp.json()
        log.debug("_call %s(%r) => %r", verb, params, result)
        return result

    @staticmethod
    def info():
        result = SslLabsCheck._call("info")
        return result

    def check(self):
        params = {"host": self.host, "publish": "off"}
        if self.status is None and self.fresh:
            # only use this flag on the initial call
            params["startNew"] = "on"
        result = self._call("analyze", params)
        self.status = result["status"]

        if self.status in ("READY", "ERROR"):
            log.debug("%s %s: %r", self.status, self.host, self.result)
            self.result = result
            return result
        else:
            log.info("Waiting on %s %s", self.host, self.status)
            return None

    def passing_grade(self):
        """returns True if grade is good enough."""
        grades = jmespath.search("endpoints[].grade", self.result)
        if not grades:
            return None
        return all(re.match(r"A[+-]?", grade) for grade in grades)


class SslLabsCheckManager:
    __slots__ = ("checks", "env")

    def __init__(self, urls: List[str], fresh: bool):
        self.checks = tuple(SslLabsCheck(url, fresh) for url in urls)
        log.debug("%s(%r)", self.__class__.__name__, self.checks)
        self.env = jinja2.Environment(
            loader=jinja2.DictLoader(_templates), undefined=jinja2.StrictUndefined
        )
        self.env.filters["style"] = click.style
        self.env.globals["manager"] = self

    def gather_reports(self):
        info = SslLabsCheck.info()
        max_assessments = int(info["maxAssessments"])
        assert max_assessments > 0
        todo = list(self.checks[::-1])
        work_slots: List[Optional[SslLabsCheck]] = [None] * max_assessments

        while True:
            # rotate work from queue through slots
            for i, check in enumerate(work_slots):
                if check is None and todo:
                    work_slots[i] = todo.pop()
                check = work_slots[i]
                if check is not None:
                    result = check.check()
                    if result:
                        if todo:
                            work_slots[i] = todo.pop()
                        else:
                            work_slots[i] = None
            # complete when nothing left to do and work slots are empty
            if todo or any(work_slots):
                time.sleep(10)
            else:
                break

        assert not any(work_slots)
        assert todo == []
        assert all(check.status in ("READY", "ERROR") and check.result for check in self.checks)

    def report_stream(self):
        template = self.env.get_template("report")
        return template.stream(manager=self)


@click.command()
@click.option("-d", "--debug", is_flag=True)
@click.option("--fresh", is_flag=True)
@click.argument("urls", nargs=-1, required=True)
def ssllabs_check(debug, fresh, urls):
    if debug:
        log.setLevel(logging.DEBUG)
    manager = SslLabsCheckManager(urls, fresh)
    manager.gather_reports()
    for line in manager.report_stream():
        click.echo(line)
    bad_domains = [check.host for check in manager.checks if not check.passing_grade()]
    if bad_domains:
        log.error("Some domains did not pass: %s", bad_domains)
        sys.exit(1)
    else:
        log.info("All domains have a passing grade")
        sys.exit(0)


if __name__ == "__main__":
    ssllabs_check()