python/testcase/shell

#!/usr/bin/env python3

"""
This command implements a shell for managing test cases.
"""

from __future__ import annotations

import argparse
import shlex
import subprocess
import tempfile
import time
import json
import os
import sys
from typing import (
    List,
    Set,
)
from urllib.parse import urlunparse
import webbrowser

import yattag.doc
import yattag.indentation

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(sys.argv[0]))))
from ooni.dataformat import (
    testcase,
)
from ooni import htmlx
import tempfiledir


#
# Browse
#
# This section implements the browse subcommand
#


def sh_browse_run(args: argparse.Namespace):
    """Generates an HTML representation of a set of testcases."""

    tcs: List[testcase.TestCase] = []
    for filename in args.file:
        if os.path.isdir(filename):
            for name in os.listdir(filename):
                fname = os.path.join(filename, name)
                if not os.path.isfile(fname):
                    continue
                if not fname.endswith(".yaml"):
                    continue
                tcs.append(testcase.load(fname))
        elif filename.endswith(".yaml"):
            tcs.append(testcase.load(filename))

    doc, _, _ = yattag.doc.Doc().tagtext()
    with doc.tag("html"):
        doc.attr(lang="en")
        with doc.tag("title"):
            doc.text("Measurex test cases")
        with doc.tag("link"):
            doc.attr(href="main.css", rel="stylesheet")
        with doc.tag("body"):
            htmlx.test_cases_overview(doc, tcs)
            for tc in tcs:
                htmlx.test_case_info(doc, tc)
            with doc.tag("script"):
                doc.attr(src="main.js")

    with tempfiledir.new_named_tempfile(dir="html", delete=False, suffix=".html") as filep:
        print("<!DOCTYPE HTML>", file=filep)
        print(yattag.indentation.indent(doc.getvalue()), file=filep)
        fullpath = os.path.abspath(filep.name)
        url = urlunparse(("file", "", fullpath, "", "", ""))
        webbrowser.open(url)


def sh_browse_configure(apa: argparse.ArgumentParser):
    """Configures the CLI parser for the browse subcommand."""
    apa.add_argument(
        "file",
        help="one or more yaml files (or a directory) to browse",
        nargs="+",
    )
    apa.set_defaults(func=sh_browse_run)


#
# Import
#
# This section implements the import subcommand
#


def sh_import_run(args: argparse.Namespace):
    """Imports one or more testcases in .tar.gz format."""
    for filename in args.file:
        if os.path.isdir(filename):
            for name in os.listdir(filename):
                fname = os.path.join(filename, name)
                if not os.path.isfile(fname):
                    continue
                if not fname.endswith(".tar.gz"):
                    continue
                testcase.import_from_tarball(args.destdir, fname)
        elif filename.endswith(".tar.gz"):
            testcase.import_from_tarball(args.destdir, filename)


def sh_import_configure(apa: argparse.ArgumentParser):
    """Configures the CLI parser for the import subcommand."""
    apa.add_argument(
        "-d",
        "--destdir",
        help="destination directory where to write imported test cases",
        required=True,
    )
    apa.add_argument(
        "file",
        help="one or more tar.gz files (or a directory) generated by the probe to import",
        nargs="+",
    )
    apa.set_defaults(func=sh_import_run)


#
# List
#
# This section implements the list command.
#


def sh_list_run(args: argparse.Namespace):
    """Lists one or more test cases."""
    tab = testcase.Tabular()
    for filename in args.file:
        if os.path.isdir(filename):
            for name in os.listdir(filename):
                fname = os.path.join(filename, name)
                if not os.path.isfile(fname):
                    continue
                if not fname.endswith(".yaml"):
                    continue
                with testcase.load(fname) as tc:
                    tab.append(tc.manifest().as_tabular())
        elif filename.endswith(".yaml"):
            with testcase.load(filename) as tc:
                tab.append(tc.manifest().as_tabular())
    print(tab.tabulatex(format=args.format, sortkey=lambda x: x[0]))


def sh_list_configure(apa: argparse.ArgumentParser):
    """Configures the CLI parser for the list subcommand."""
    apa.add_argument(
        "file",
        help="one or more test cases (or a directory) to list",
        nargs="+",
    )
    apa.add_argument(
        "--format",
        help="choose the format (e.g., grid, github, html, ...)",
        default="grid",
    )
    apa.set_defaults(func=sh_list_run)


#
# Show
#
# This section contains code to show a single test case.
#


def sh_show_run(args: argparse.Namespace):
    """Shows one or more test cases."""
    for filename in args.file:
        if not os.path.isfile(filename):
            continue
        if not filename.endswith(".yaml"):
            continue
        with testcase.load(filename) as tc:
            cache = tc.cache()
            dns = cache.dns_measurements_as_tabular()
            print(dns.tabulatex(format=args.format))
            epnt = cache.endpoint_measurements_as_tabular()
            print(epnt.tabulatex(format=args.format, sortkey=lambda x: x[4]))


def sh_show_configure(apa: argparse.ArgumentParser):
    """Configures the CLI parser for the show subcommand."""
    apa.add_argument(
        "file",
        help="test case in yaml format to show",
        nargs=1,
    )
    apa.add_argument(
        "--format",
        help="choose the format (e.g., grid, github, html, ...)",
        default="grid",
    )
    apa.set_defaults(func=sh_show_run)


#
# Decode
#
# This section decodes DNS and HTTP round trips.
#


def sh_decode_idfilter(args: argparse.Namespace) -> Set[int]:
    """Returns the correct ID filter for decode."""
    out: Set[int] = set()
    if args.id is not None:
        out = set([int(x) for x in args.id])
    return out


def sh_decode_run(args: argparse.Namespace):
    """Decodes DNS and HTTP round trips."""
    idfilter = sh_decode_idfilter(args)
    for filename in args.file:
        if not os.path.isfile(filename):
            continue
        if not filename.endswith(".yaml"):
            continue
        with testcase.load(filename) as tc:
            cache = tc.cache()
            for dns in cache.dns_measurements(idfilter):
                print(dns.decode())
            for epnt in cache.endpoint_measurements(idfilter):
                print(epnt.decode())


def sh_decode_configure(apa: argparse.ArgumentParser):
    """Configures the CLI parser for the decode subcommand."""
    apa.add_argument(
        "--id",
        help="limit decoding to specified ID (may be specified multiple times)",
        action="append",
    )
    apa.add_argument(
        "file",
        help="test case in yaml format to decode",
        nargs=1,
    )
    apa.set_defaults(func=sh_decode_run)


#
# Rerun
#
# This section contains code to rerun a given test case
#


def sh_rerun_write_cache_into(basedir: str, cache: testcase.Cache):
    """Writes the cache into the given basedir."""
    for key, value in cache.raw_entries():
        dirname = os.path.join(basedir, os.path.dirname(key))
        os.makedirs(dirname, exist_ok=True)
        fullname = os.path.join(basedir, key)
        with open(fullname, "w") as outfp:
            json.dump(value, outfp)


def sh_rerun_start_thd(basedir: str):
    """Starts the test helper daemon."""
    argv = [
        "./thd",
        "--cache-forever",
        "-C",
        os.path.join(basedir, "testcase", "cache", "th"),
        "-L",
        os.path.join(basedir, "thd-log.txt"),
        "-Nv",
    ]
    print(f"about to start thd: {shlex.join(argv)}")
    proc = subprocess.Popen(argv)
    time.sleep(2)
    return proc


def sh_rerun_run_probe(emoji: bool, basedir: str, target_url: str):
    """Runs the probe."""
    argv = [
        "./websteps",
        "-b",
        "ws://127.0.0.1:9876/websteps/v1/websocket",
        "-C",
        os.path.join(basedir, "testcase", "cache", "probe"),
        "-i",
        target_url,
        "-L",
        os.path.join(basedir, "probe-log.txt"),
        "-No",
        os.path.join(basedir, "report.jsonl"),
        "-Pv",
    ]
    if emoji:
        argv.append("-e")
    print(f"about to run probe: {shlex.join(argv)}")
    subprocess.run(argv, check=True)


def sh_rerun_run(args: argparse.Namespace):
    """Reruns the given test case."""
    for filename in args.file:
        if not os.path.isfile(filename):
            continue
        if not filename.endswith(".yaml"):
            continue
        with testcase.load(filename) as tc:
            manifest = tc.manifest()
            cache = tc.cache()
            basedir = tempfiledir.new_tempdir()
            sh_rerun_write_cache_into(basedir, cache)
            target_url = manifest.url
            thd = sh_rerun_start_thd(basedir)
            sh_rerun_run_probe(args.emoji, basedir, target_url)
            print(f"output available in {basedir}")
            thd.terminate()
            print(f"waiting for {thd} to terminate...")
            thd.wait()


def sh_rerun_configure(apa: argparse.ArgumentParser):
    """Configures the CLI parser for the rerun subcommand."""
    apa.add_argument(
        "-e",
        "--emoji",
        help="enables emoji when executing websteps",
        action="store_true",
    )
    apa.add_argument(
        "file",
        help="test case in yaml format to rerun",
        nargs=1,
    )
    apa.set_defaults(func=sh_rerun_run)


#
# Main
#
# This section contains the main functionality
#


def getopt():
    """Parse command line arguments."""
    cli = argparse.ArgumentParser(
        prog="./python/testcase/shell", description="shell for managing test cases"
    )
    subcommands = cli.add_subparsers()

    cmd_browse = subcommands.add_parser(
        "browse",
        help="makes a set of test cases browseble",
    )
    sh_browse_configure(cmd_browse)

    cmd_import = subcommands.add_parser(
        "import",
        help="imports test cases from .tar.gz files generated by the probe",
    )
    sh_import_configure(cmd_import)

    cmd_list = subcommands.add_parser(
        "list", help="lists test cases in yaml format", aliases=["ls"]
    )
    sh_list_configure(cmd_list)

    cmd_show = subcommands.add_parser("show", help="shows the content of test cases")
    sh_show_configure(cmd_show)

    cmd_decode = subcommands.add_parser(
        "decode", help="decodes DNS and HTTP round trips"
    )
    sh_decode_configure(cmd_decode)

    cmd_rerun = subcommands.add_parser(
        "rerun",
        help="reruns a given test case",
    )
    sh_rerun_configure(cmd_rerun)

    return cli.parse_args(), cli


def main():
    parsed_args, cli = getopt()
    if not hasattr(parsed_args, "func"):
        cli.print_help()
        sys.exit(0)
    parsed_args.func(parsed_args)


if __name__ == "__main__":
    main()