import json
import logging
import itertools
import functools
from pathlib import Path
import multiprocessing
import multiprocessing.pool
from typing import Optional

import numpy as np
import pandas as pd
from lab import tracev2
from lab.defences import tamaraw

import common
from common.doceasy import doceasy, Use, Or

_LOGGER = logging.getLogger(__name__)


def main(
    input_dir: Path,
    defence: str,
    tamaraw_config: Optional[str],
    n_jobs: int
):
    """Calculate the overhead using multiple processes."""
    common.init_logging()
    _LOGGER.info("Using parameters: %s", locals())

    if defence == "tamaraw":
        if tamaraw_config is None:
            raise ValueError("Tamaraw configuration required.")
        tamaraw_config = json.loads(tamaraw_config)

    assert input_dir.is_dir(), f"invalid path {input_dir}"
    directories = sorted(
        [x.parent for x in Path(input_dir).glob("**/defended/")]
    )
    _LOGGER.info("Found %d samples", len(directories))

    func = functools.partial(
        _calculate_overhead, defence=defence, tamaraw_config=tamaraw_config)
    if n_jobs > 1:
        chunksize = max(len(directories) // (n_jobs * 2), 1)
        with multiprocessing.pool.Pool(n_jobs) as pool:
            scores = list(
                pool.imap_unordered(func, directories, chunksize=chunksize)
            )
    else:
        # Run in the main process
        scores = [func(x) for x in directories]
    _LOGGER.info("Overhead calculation complete")

    results = pd.DataFrame.from_records(itertools.chain.from_iterable(scores))
    print(results.to_csv(header=True, index=False), end="")


def _parse_duration(path):
    """Return the time taken to download all of the application's HTTP
    resources in ms, irrespective of any additional padding or traffic,
    as logged by the run.
    """
    tag = "[FlowShaper] Application complete after "  # xxx ms
    found = None
    with (path / "stdout.txt").open(mode="r") as stdout:
        found = [line for line in stdout if line.startswith(tag)][-1]
    assert found, f"Run never completed! {path}"

    # Parse the next word as an integer
    return int(found[len(tag):].split()[0])


def _calculate_overhead(dir_, *, defence: str, tamaraw_config):
    try:
        control = np.sort(tracev2.from_csv(dir_ / "undefended" / "trace.csv"))
        defended = np.sort(tracev2.from_csv(dir_ / "defended" / "trace.csv"))
        schedule = np.sort(tracev2.from_csv(dir_ / "defended" / "schedule.csv"))
    except Exception as err:
        raise ValueError(f"Error loading files in {dir_}") from err

    undefended_size = np.sum(np.abs(control["size"]))
    defended_size = np.sum(np.abs(defended["size"]))
    simulated_size = np.sum(np.abs(schedule["size"]))
    simulated_size_alt = None
    # Add the undefended size as padding only defences only list the padding
    # in the schedule.
    if defence == "front":
        simulated_size += undefended_size

    if defence == "tamaraw":
        tamaraw_trace = tamaraw.simulate(
            control,
            packet_size=tamaraw_config["packet_size"],
            rate_in=tamaraw_config["rate_in"] / 1000,
            rate_out=tamaraw_config["rate_out"] / 1000,
            pad_multiple=tamaraw_config["packet_multiple"],
        )
        simulated_size_alt = np.sum(np.abs(tamaraw_trace["size"]))

    assert control["time"][0] == 0, "trace must start at 0s"
    # The trace should also already be sorted and in seconds
    undefended_ms = int(control["time"][-1] * 1000)
    defended_ms = _parse_duration(dir_ / "defended")
    if defence == "front":
        simulated_ms = undefended_ms
    elif defence == "tamaraw":
        unpadded_tamaraw = tamaraw.simulate(
            control,
            packet_size=tamaraw_config["packet_size"],
            rate_in=tamaraw_config["rate_in"] / 1000,
            rate_out=tamaraw_config["rate_out"] / 1000,
            pad_multiple=1,
        )
        simulated_ms = int(unpadded_tamaraw["time"][-1] * 1000)
    else:
        raise ValueError(f"Unsupported defence: {defence}")

    return [
        {
            "sample": str(dir_),
            "overhead": "bandwidth",
            "setting": "collected",
            "value": (defended_size - undefended_size) / undefended_size
        },
        {
            "sample": str(dir_),
            "overhead": "bandwidth",
            "setting": "simulated",
            "value": (simulated_size - undefended_size) / undefended_size
        },
        {
            "sample": str(dir_),
            "overhead": "bandwidth",
            "setting": "simulated-alt",
            "value": ((simulated_size_alt - undefended_size) / undefended_size
                      if simulated_size_alt is not None else None)
        },
        {
            "sample": str(dir_),
            "overhead": "latency",
            "setting": "collected",
            "value": (defended_ms - undefended_ms) / undefended_ms
        },
        {
            "sample": str(dir_),
            "overhead": "latency",
            "setting": "simulated",
            "value": (simulated_ms - undefended_ms) / undefended_ms
        },
    ]


if __name__ == "__main__":
    main(**doceasy(__doc__, {
        "DEFENCE": Or("tamaraw", "front"),
        "INPUT_DIR": Use(Path),
        "--n-jobs": Use(int),
        "--tamaraw-config": Or(None, str),
    }))

import logging
import itertools
import functools
from pathlib import Path
import multiprocessing
import multiprocessing.pool
from typing import Sequence, Dict, Optional

import numpy as np
import pandas as pd
import tslearn.metrics
from scipy.stats import pearsonr
from scipy.spatial.distance import euclidean

import common
from common import timeseries

_LOGGER = logging.getLogger(__name__)


def main(
    input_,
    output,
    *,
    defence: str,
    ts_offset: Dict[str, int],
    resample_rates: Sequence["str"],
    lcss_eps: int,
    filter_below: Sequence[int] = (0,),
    jobs: Optional[int] = None,
):
    """Score how close a defended time series is to the theoretical."""
    common.init_logging()
    _LOGGER.info("Using parameters: %s", locals())

    directories = sorted([x.parent for x in Path(input_).glob("**/defended/")])
    _LOGGER.info("Found %d samples", len(directories))

    jobs = jobs or (multiprocessing.cpu_count() or 4)
    func = functools.partial(
        _calculate_score,
        defence=defence,
        ts_offset=ts_offset,
        resample_rates=resample_rates,
        filter_below=filter_below,
        lcss_eps=lcss_eps,
    )

    if jobs > 1:
        chunksize = max(len(directories) // (jobs * 2), 1)
        with multiprocessing.pool.Pool(jobs) as pool:
            scores = list(pool.imap_unordered(func, directories, chunksize=chunksize))
    else:
        # Run in the main process
        scores = list(map(func, directories))
    _LOGGER.info("Score calculation complete")

    pd.DataFrame.from_records(itertools.chain.from_iterable(scores)).to_csv(
        output, header=True, index=False
    )


def _calculate_score(
    dir_,
    *,
    defence: str,
    ts_offset,
    resample_rates,
    filter_below,
    lcss_eps,
):
    """Score how close a padding-only defended time series is to
    padding schedule.
    """
    assert defence in ("front", "tamaraw")
    pad_only = defence == "front"

    schedule_ts = timeseries.from_csv(dir_ / "defended" / "schedule.csv")
    defended_ts = timeseries.from_csv(dir_ / "defended" / "trace.csv")
    control_ts = timeseries.from_csv(dir_ / "undefended" / "trace.csv")

    offsets = range(ts_offset["min"], ts_offset["max"], ts_offset["inc"])
    simulated_ts = simulate_with_lag(
        control_ts, schedule_ts, defended_ts, offsets, pad_only=pad_only
    )

    results = []
    for rate, direction, min_pkt_size in itertools.product(
        resample_rates, ("in", "out"), filter_below
    ):
        series = pd.DataFrame(
            {
                "a": timeseries.resample(
                    _filter(defended_ts[direction], min_pkt_size), rate
                ),
                "b": timeseries.resample(
                    _filter(simulated_ts[direction], min_pkt_size), rate
                ),
                "c": timeseries.resample(
                    _filter(control_ts[direction], min_pkt_size), rate
                ),
            }
        ).fillna(0)

        _LOGGER.debug("Series length at rate %s: %d", rate, len(series))
        _LOGGER.debug("Series summary at rate %s: %s", rate, series.describe())

        if len(series["a"]) > 30_000:
            _LOGGER.warning(
                "Refusing to allocate >= 8 GB for sample: %s, %s, %s",
                dir_,
                rate,
                direction,
            )
            continue

        results.append(
            {
                "sample": str(dir_),  # Sample name
                "rate": rate,
                "dir": direction,
                "min_pkt_size": min_pkt_size,
                "pearsonr": pearsonr(series["a"], series["b"])[0],
                "lcss": tslearn.metrics.lcss(series["a"], series["b"], eps=lcss_eps),
            }
        )
    return results


def _filter(column, below):
    assert column.min() >= 0
    return column[column >= below]


def simulate_with_lag(
    control,
    schedule,
    defended,
    offsets,
    pad_only: bool,
    rate="5ms",
):
    """Find the best offset such that the simulated trace formed by
    combining the control and schedule where the schedule is lagged
    the offset, has the lowest euclidean distance to the defended trace.

    When pad_only is False, the schedule is taken as the already
    simulated trace
    """
    defended_in = timeseries.resample(defended["in"], rate)
    simulated_out = (
        control["out"].append(schedule["out"]) if pad_only else schedule["out"]
    )

    best_offset = 0
    best_distance = np.inf
    best_simulated = None

    for offset in offsets:
        shifted_schedule = pd.Series(
            schedule["in"].values,
            index=(schedule["in"].index + pd.Timedelta(f"{offset}ms")),
        )
        simulated_in = (
            control["in"].append(shifted_schedule) if pad_only else shifted_schedule
        )

        # Put together in a dataframe to ensure they have the same length and
        # indices
        frame = pd.DataFrame(
            {
                "defended": defended_in,
                "simulated": timeseries.resample(simulated_in, rate),
            }
        ).fillna(0)

        distance = euclidean(frame["defended"], frame["simulated"])
        if distance < best_distance:
            best_offset = offset
            best_distance = distance
            best_simulated = simulated_in

    assert best_simulated is not None

    _LOGGER.debug("Using an incoming offset of %d ms", best_offset)
    return pd.DataFrame(
        {
            # Take the sum of any given time instance to handle rare duplicates
            "in": best_simulated.groupby("time").sum(),
            "out": simulated_out.groupby("time").sum(),
        }
    ).fillna(0)


if __name__ == "__main__":
    main(
        str(snakemake.input[0]),
        str(snakemake.output[0]),
        **snakemake.params,
        jobs=snakemake.threads,
    )

import logging
from pathlib import Path
from itertools import product, islice
from typing import Optional, Dict
import h5py
import numpy as np
import lab.tracev2 as trace
from lab.defences import front
import common

LABELS_DTYPE = np.dtype([("class", "i4"), ("region", "i4"), ("sample", "i4")])


def main(
    input_: str,
    output: str,
    *,
    n_monitored: int,
    n_instances: int,
    n_unmonitored: int,
    n_regions: int,
    simulate: str = "",
    simulate_kws: Dict = None,
):
    """Create an HDF5 dataset at output using the files found in the
    directory input_.  The CSV traces should be found under

        <input_>/<page_id>/<region>_<sample>/

    """
    common.init_logging()
    rng = None
    if simulate_kws and "seed" in simulate_kws:
        rng = np.random.default_rng(simulate_kws["seed"])
        del simulate_kws["seed"]
    if simulate_kws and "use_empty_resources" in simulate_kws:
        del simulate_kws["use_empty_resources"]

    sample_dirs = [p for p in Path(input_).iterdir() if p.is_dir()]
    # Sort such that the dirs with the most samples are first ties broken by
    # lower sample ids
    sample_dirs.sort(key=lambda p: (sum(-1 for _ in p.glob("*")), int(p.stem)))

    if len(sample_dirs) < n_monitored + n_unmonitored:
        raise ValueError(f"Insufficient samples: {len(sample_dirs)}")
    sample_dirs = sample_dirs[:(n_monitored + n_unmonitored)]

    n_rows = n_monitored * n_instances + n_unmonitored
    labels = np.recarray(n_rows, dtype=LABELS_DTYPE)
    sizes = np.full(n_rows, None, dtype=object)
    timestamps = np.full(n_rows, None, dtype=object)

    index = 0
    for label, sample_path in enumerate(sample_dirs[:n_monitored]):
        # Tterator that takes one sample from each region before taking
        # the second sample from a region. This is akin to how we
        # generated the samples.
        sample_iter = product(range(n_instances), range(n_regions))

        # Take only the first n_instances from the infinite iterator
        for (id_, region_id) in islice(sample_iter, n_instances):
            index = (label * n_instances) + (id_ * n_regions) + region_id
            labels[index] = (label, region_id, id_)

            sample = _maybe_simulate(
                sample_path / f"{region_id}_{id_}", simulate, simulate_kws, rng
            )
            sizes[index] = sample["size"]
            timestamps[index] = sample["time"]

    region_index = 0
    for sample_path in sample_dirs[n_monitored:(n_monitored + n_unmonitored)]:
        paths = [(region_id, sample_path / f"{region_id}_0")
                 for region_id in range(n_regions)]
        paths = [(region_id, p) for (region_id, p) in paths if p.exists()]
        assert paths, f"no samples in directory? {sample_path}"

        # Select a path from the list, cycling among the used region if there
        # are more than one regions, otherwise returning an inde in the paths
        # list
        region_id, path = paths[(region_index % n_regions) % len(paths)]
        region_index += 1

        index = index + 1
        labels[index] = (-1, region_id, 0)

        sample = _maybe_simulate(path, simulate, simulate_kws, rng)
        sizes[index] = sample["size"]
        timestamps[index] = sample["time"]
    assert index == (n_rows - 1), "not enough samples"

    order = labels.argsort(order=["class", "region"])
    with h5py.File(output, mode="w") as h5out:
        h5out.create_dataset("/labels", dtype=LABELS_DTYPE, data=labels[order])
        h5out.create_dataset(
            "/sizes", dtype=h5py.vlen_dtype(np.dtype("i4")), data=sizes[order]
        )
        h5out.create_dataset(
            "/timestamps", dtype=h5py.vlen_dtype(np.dtype(float)),
            data=timestamps[order]
        )


def _maybe_simulate(path: Path, simulate: str, simulate_kws, rng):
    try:
        if not simulate:
            return trace.from_csv(path / "trace.csv")

        if simulate == "tamaraw":
            return trace.from_csv(path / "schedule.csv")
        if simulate == "front":
            simulate_kws = {
                key: value for (key, value) in simulate_kws.items()
                if key in [
                    "max_client_packets", "max_server_packets", "packet_size",
                    "peak_minimum", "peak_maximum", "random_state"
                ]
            }
            assert rng is not None
            assert "seed" not in simulate_kws
            assert all(kw in simulate_kws for kw in [
                "max_client_packets", "max_server_packets", "packet_size",
                "peak_minimum", "peak_maximum",
            ])

            padding = front.generate_padding(**simulate_kws, random_state=rng)
            # We should be pointing to an undefended trace.csv file
            baseline = trace.from_csv(path / "trace.csv")
            return front.simulate(baseline, padding)
        raise ValueError(f"Unrecognised simulation: {simulate}")
    except Exception:
        logging.error("Failed on sample: %s", path)
        raise


if __name__ == "__main__":
    snakemake = globals().get("snakemake", None)
    main(
        input_=str(snakemake.input[0]),
        output=str(snakemake.output[0]),
        n_regions=snakemake.config["wireguard"]["n_regions"],
        **snakemake.params
    )

import time
import logging
import dataclasses
from pathlib import Path
from typing import Optional, ClassVar, Sequence, Union

import h5py
import numpy as np
from sklearn.metrics import make_scorer
from sklearn.model_selection import GridSearchCV
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import FunctionTransformer
from sklearn.model_selection import StratifiedKFold, train_test_split
from lab.feature_extraction.trace import ensure_non_ragged
from lab.classifiers import dfnet
from lab.metrics import rprecision_score, recall_score
import tensorflow

from common import doceasy
from common.doceasy import Use, Or, And

PAPER_EPOCHS: int = 30
MAX_RAND_SEED: int = 10_000

# Hyperparameters from the paper
DEFAULT_N_PACKETS: int = 5000
DEFAULT_LEARNING_RATE: float = 0.002
DEFAULT_EPOCHS: int = 30


def main(outfile, **kwargs):
    """Create and run the experiment and write the CSV results to outfile."""
    logging.basicConfig(
        format='[%(asctime)s] [%(levelname)s] %(name)s - %(message)s',
        level=logging.INFO
    )
    # Use autoclustering for speedup
    tensorflow.config.optimizer.set_jit("autoclustering")

    (probabilities, y_true, classes) = Experiment(**kwargs).run()

    y_pred = classes[np.argmax(probabilities, axis=1)]
    score = rf1_score(y_true, y_pred)
    logging.info("r_20 f1-score = %.4g", score)

    outfile.writerow(["y_true"] + list(classes))
    outfile.writerows(
        np.hstack((np.reshape(y_true, (-1, 1)), probabilities)))


@dataclasses.dataclass
class Experiment:
    """An experiment to evalute hyperparameter tuned DF classifier.
    """
    # The path to the dataset of sizes and times
    dataset_path: Path

    # The fraction of samples to use for testing the final model
    test_size: float = 0.1

    # Number of folds used in the stratified k-fold cross validation
    n_folds: int = 3

    # Level of debugging output from sklearn and tensorflow
    verbose: int = 0

    # Random seed for the experiment
    seed: int = 114155

    # Hyperparams to use if not "tune"
    hyperparams: Union[str, dict] = "tune"

    # Hyperparameters to search
    n_packet_parameters: Sequence[int] = (5_000, 7_500, 10_000)
    tuned_parameters: dict = dataclasses.field(default_factory=lambda: {
        "epochs": [30],
        "learning_rate": [0.002],
    })

    # Other seeds which are chosen for different operations
    seeds_: Optional[dict] = None

    logger: ClassVar = logging.getLogger("Experiment")

    def run(self):
        """Run hyperparameter tuning for the DeepFingerprinting classifier
        and return the prediction probabilities for the best chosen
        classifier.
        """
        # Generate and set random seeds
        rng = np.random.default_rng(self.seed)
        self.seeds_ = {
            "train_test": rng.integers(MAX_RAND_SEED),
            "kfold_shuffle": rng.integers(MAX_RAND_SEED),
            "tensorflow": rng.integers(MAX_RAND_SEED),
        }
        tensorflow.random.set_seed(self.seeds_["tensorflow"])

        self.logger.info("Running %s", self)
        start = time.perf_counter()

        # Load the dataset
        X, y = self.load_dataset()
        n_classes = len(np.unique(y))
        self.logger.info("Dataset shape=%s, n_classes=%d", X.shape, n_classes)

        # Generate our training and final testing set
        x_train, x_test, y_train, y_test = train_test_split(
            X, y, test_size=self.test_size, stratify=y, shuffle=True,
            random_state=self.seeds_["train_test"]
        )

        if self.hyperparams == "tune":
            self.logger.info("Performing hyperparameter tuning ...")
            # Tune other hyperparameters and fit the final estimator
            classifier = self.tune_hyperparameters(
                x_train, y_train, n_classes=n_classes
            )
        else:
            assert isinstance(self.hyperparams, dict)
            n_packets = self.hyperparams.get("n_packets", DEFAULT_N_PACKETS)
            learning_rate = self.hyperparams.get(
                "learning_rate", DEFAULT_LEARNING_RATE
            )
            epochs = self.hyperparams.get("epochs", DEFAULT_EPOCHS)
            self.logger.info(
                "Using n_packets=%s, learning_rate=%.3g, and epochs=%d",
                n_packets, learning_rate, epochs
            )

            x_train = first_n_packets(x_train, n_packets=n_packets)
            x_test = first_n_packets(x_test, n_packets=n_packets)

            classifier = dfnet.DeepFingerprintingClassifier(
                n_classes=n_classes, verbose=min(self.verbose, 1),
                n_features=n_packets, epochs=epochs, learning_rate=learning_rate
            )
            classifier.fit(x_train, y_train)

        # Predict the classes for the test set
        probabilities = classifier.predict_proba(x_test)
        self.logger.info(
            "Experiment complete in %.2fs.", (time.perf_counter() - start))

        return (probabilities, y_test, classifier.classes_)

    def tune_hyperparameters(self, x_train, y_train, *, n_classes):
        """Perform hyperparameter tuning on the learning rate."""
        assert self.seeds_ is not None, "seeds must be set"
        pipeline = Pipeline([
            ("first_n_packets", FunctionTransformer(first_n_packets)),
            ("dfnet", dfnet.DeepFingerprintingClassifier(
                n_classes=n_classes, verbose=min(self.verbose, 1),
            ))
        ])
        param_grid = [
            {
                "first_n_packets__kw_args": [{"n_packets": n_packets}],
                "dfnet__n_features": [n_packets],
                **{
                    f"dfnet__{key}": values
                    for key, values in self.tuned_parameters.items()
                }
            }
            for n_packets in self.n_packet_parameters
        ]
        cross_validation = StratifiedKFold(
            self.n_folds, shuffle=True,
            random_state=self.seeds_["kfold_shuffle"]
        )

        grid_search = GridSearchCV(
            pipeline, param_grid, cv=cross_validation, error_score="raise",
            scoring=make_scorer(rf1_score), verbose=self.verbose, refit=True,
        )
        grid_search.fit(x_train, y_train)

        self.logger.info("hyperparameter results = %s", grid_search.cv_results_)
        self.logger.info("hyperparameter best = %s", grid_search.best_params_)

        return grid_search.best_estimator_

    def load_dataset(self):
        """Load the features and classes from the dataset. Slices the
        packet features to the maximum evaluated.
        """
        max_packets = max(self.n_packet_parameters)
        with h5py.File(self.dataset_path, mode="r") as h5in:
            features = ensure_non_ragged(h5in["sizes"], dimension=max_packets)
            classes = np.asarray(h5in["labels"]["class"][:])
        return features, classes


def first_n_packets(features, *, n_packets: int):
    """Return the first n_packets packets along with the meta features."""
    return features[:, :n_packets]


def rf1_score(y_true, y_pred, *, negative_class=-1, ratio=20):
    """Compute the F1-score using the r-precisions with the specified ratio
    and recall.
    """
    precision = rprecision_score(
        y_true, y_pred, negative_class=negative_class, ratio=ratio,
        # If we're dividing by zero it means there were no true positives and
        # thus recall will be zero and the F1 score below will be zero.
        zero_division=1
    )
    recall = recall_score(y_true, y_pred, negative_class=negative_class)
    return 2 * precision * recall / (precision + recall)


if __name__ == "__main__":
    main(**doceasy.doceasy(__doc__, {
        "OUTFILE": doceasy.CsvFile(mode="w", default="-"),
        "DATASET_PATH": Use(Path),
        "--verbose": Use(int),
        "--hyperparams": Or("tune", And(doceasy.Mapping(), {
            doceasy.Optional("n_packets"): Use(int),
            doceasy.Optional("epochs"): Use(int),
            doceasy.Optional("learning_rate"): Use(float),
        }))
    }))