import json, os, sys, hmac, hashlib, datetime
import urllib.parse, urllib.request, urllib.error
import xml.etree.ElementTree as ET


def _cfg():
    return json.loads(os.environ.get("INTEGRATION_SECRETS", "{}"))


def _inputs():
    return json.loads(os.environ.get("INTEGRATION_INPUTS", "{}"))


def _sign_key(key, date_stamp, region, service):
    def _h(k, m):
        return hmac.new(k, m.encode("utf-8"), hashlib.sha256).digest()
    k_date = _h(("AWS4" + key).encode("utf-8"), date_stamp)
    k_region = _h(k_date, region)
    k_service = _h(k_region, service)
    return _h(k_service, "aws4_request")


def _strip_ns(tag):
    return tag.split("}", 1)[1] if "}" in tag else tag


def _xml_to_dict(elem):
    d = {}
    children = list(elem)
    if not children:
        return (elem.text or "").strip()
    for c in children:
        tag = _strip_ns(c.tag)
        val = _xml_to_dict(c)
        if tag in d:
            if not isinstance(d[tag], list):
                d[tag] = [d[tag]]
            d[tag].append(val)
        else:
            d[tag] = val
    return d


def aws_query(service, host, region, action, version, params, cfg):
    # params: dict of extra query params for this Action
    body_params = {"Action": action, "Version": version}
    body_params.update({k: str(v) for k, v in params.items() if v is not None})
    body = urllib.parse.urlencode(sorted(body_params.items()))

    access_key = str(cfg.get("access_key_id", ""))
    secret_key = str(cfg.get("secret_access_key", ""))
    session_token = cfg.get("session_token") or ""

    now = datetime.datetime.utcnow()
    amz_date = now.strftime("%Y%m%dT%H%M%SZ")
    date_stamp = now.strftime("%Y%m%d")

    method = "POST"
    canonical_uri = "/"
    canonical_querystring = ""
    payload_hash = hashlib.sha256(body.encode("utf-8")).hexdigest()
    canonical_headers = "content-type:application/x-www-form-urlencoded; charset=utf-8\n" \
                        "host:" + host + "\n" \
                        "x-amz-date:" + amz_date + "\n"
    signed_headers = "content-type;host;x-amz-date"
    if session_token:
        canonical_headers += "x-amz-security-token:" + session_token + "\n"
        signed_headers = "content-type;host;x-amz-date;x-amz-security-token"
    canonical_request = "\n".join([method, canonical_uri, canonical_querystring,
                                   canonical_headers, signed_headers, payload_hash])

    algorithm = "AWS4-HMAC-SHA256"
    credential_scope = date_stamp + "/" + region + "/" + service + "/aws4_request"
    string_to_sign = "\n".join([algorithm, amz_date, credential_scope,
                                hashlib.sha256(canonical_request.encode("utf-8")).hexdigest()])
    signing_key = _sign_key(secret_key, date_stamp, region, service)
    signature = hmac.new(signing_key, string_to_sign.encode("utf-8"), hashlib.sha256).hexdigest()

    authorization = (algorithm + " Credential=" + access_key + "/" + credential_scope +
                     ", SignedHeaders=" + signed_headers + ", Signature=" + signature)
    headers = {
        "Content-Type": "application/x-www-form-urlencoded; charset=utf-8",
        "X-Amz-Date": amz_date,
        "Authorization": authorization,
        "Accept": "application/json",
    }
    if session_token:
        headers["X-Amz-Security-Token"] = session_token

    req = urllib.request.Request("https://" + host + "/", data=body.encode("utf-8"), headers=headers, method="POST")
    with urllib.request.urlopen(req, timeout=60) as r:
        raw = r.read()
    root = ET.fromstring(raw)
    return {_strip_ns(root.tag): _xml_to_dict(root)}


def _region(cfg, inputs):
    return (inputs.get("region") if inputs else None) or cfg.get("region") or "us-east-1"


def ec2(action, params, cfg, inputs):
    region = _region(cfg, inputs)
    return aws_query("ec2", "ec2." + region + ".amazonaws.com", region, action, "2016-11-15", params, cfg)


def iam(action, params, cfg):
    return aws_query("iam", "iam.amazonaws.com", "us-east-1", action, "2010-05-08", params, cfg)


def sts(action, params, cfg):
    return aws_query("sts", "sts.amazonaws.com", "us-east-1", action, "2011-06-15", params, cfg)


def _run(fn):
    try:
        print(json.dumps(fn(_cfg(), _inputs())))
    except urllib.error.HTTPError as e:
        print(json.dumps({"error": "HTTP " + str(e.code), "detail": e.read().decode("utf-8", "replace")}))
        sys.exit(1)
    except Exception as e:
        print(json.dumps({"error": str(e)}))
        sys.exit(1)


def main(cfg, inputs):
    identity = sts("GetCallerIdentity", {}, cfg)
    return {"ok": True, "identity": identity}


_run(main)
