#!/opt/conda/conda-bld/dxpy_1776118823485/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehol/bin/python
#
# Copyright (C) 2013-2024 DNAnexus, Inc.
#
# This file is part of dx-toolkit (DNAnexus platform client libraries).
#
#   Licensed under the Apache License, Version 2.0 (the "License"); you may not
#   use this file except in compliance with the License. You may obtain a copy
#   of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#   WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#   License for the specific language governing permissions and limitations
#   under the License.

import dxpy
import os, argparse, sys


def _parse_args():
    """
    Parse the input arguments.
    """
    parser = argparse.ArgumentParser(
        description="calls job-xxxx/getIdentityToken and retrieves a JWT token based on aud and subject claims input"
    )
    parser.add_argument(
        "--aud", help="Audience URI the JWT is intended for", required=True
    )
    parser.add_argument(
        "--subject_claims",
        action="append",
        metavar="<subject_claims>",
        help="Defines the subject claims to be validated by the cloud provider",
    )

    return parser.parse_args()


def main(aud, subject_claims):
    job_get_identity_token_input = {}

    # Parse out audience and subject_claims from args
    job_get_identity_token_input["audience"] = aud

    if subject_claims is not None:
        # Iterate over subject_claims and flatten them into a single list
        subject_claims_input = []
        for subject_claim in subject_claims:
            subject_claims_input.extend(subject_claim.split(","))
        job_get_identity_token_input["subject_claims"] = subject_claims_input

    # Call job-xxxx/getIdentityToken
    if "DX_JOB_ID" in os.environ:
        response = dxpy.api.job_get_identity_token(
            dxpy.JOB_ID, job_get_identity_token_input
        )
        sys.stdout.write(response["Token"])
    else:
        print("This script should be run from within a job environment.")
        sys.exit(1)


if __name__ == "__main__":
    args = _parse_args()
    main(args.aud, args.subject_claims)
