"""mitmproxy addon: run an agent (or a whole container) through a proxy that enforces the policy.

    mitmdump -s "$(agent-egress-guard mitm-script)"

Environment variables: EGRESS_GUARD_RULES, EGRESS_GUARD_HOSTS (licensed rule files),
EGRESS_GUARD_API_KEY (page-type database), EGRESS_GUARD_STRICT=1.
Denied requests get a 403 response with the verdict as JSON.
"""
import json
import os

from mitmproxy import http  # noqa: F401  (only available inside mitmproxy)

from agent_egress_guard.engine import Guard


class EgressGuardAddon:
    def __init__(self):
        self.guard = Guard(rules_path=os.environ.get("EGRESS_GUARD_RULES") or None,
                           hosts_path=os.environ.get("EGRESS_GUARD_HOSTS") or None,
                           api_key=os.environ.get("EGRESS_GUARD_API_KEY") or None,
                           strict=os.environ.get("EGRESS_GUARD_STRICT") == "1")

    def request(self, flow):
        v = self.guard.check(flow.request.pretty_url, flow.request.method)
        if v.decision == "deny":
            flow.response = http.Response.make(403, json.dumps({"egress_guard": v.as_dict()}),
                                               {"Content-Type": "application/json"})
        elif v.decision == "flag":
            flow.request.headers["X-Egress-Guard"] = "flag:" + v.rule


addons = [EgressGuardAddon()]