Google Key Management Service

import demistomock as demisto
from CommonServerPython import *

from CommonServerUserPython import *

"""IMPORTS"""
import base64
from json import JSONDecodeError
from typing import Any

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding
from google.cloud import kms  # type: ignore[attr-defined]

"""
For further information about the API used in the integration see:

1) Google KMS API Client libraries information:
    https://cloud.google.com/kms/docs/reference/libraries

2) Git resource with some API use examples:
    https://github.com/GoogleCloudPlatform/python-docs-samples/tree/master/kms/api-client

"""

INTEGRATION_NAME = "Google Key Management System"
# lowercase with `-` dividers
INTEGRATION_COMMAND_NAME = "google-kms"
# No dividers
INTEGRATION_CONTEXT_NAME = "GoogleKMS"

DEMISTO_DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%S"
RFC3339_DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"


class Client:
    def __init__(self, params: dict[str, Any]):
        self.project = params.get("project")
        self.location = params.get("location")
        self.key_ring = params.get("key_ring")
        self.service_account = params.get("credentials_service_account", {}).get("password") or params.get("service_account")
        if not self.service_account:
            raise DemistoException("User's Service Account JSON must be provided.")
        handle_proxy()
        # Creates an API client for the KMS API.
        try:
            self.kms_client = self._init_kms_client()

        except JSONDecodeError:
            raise Exception("Service Account json has missing details. You need to re-create the json file.")

    def _init_kms_client(self):
        """Creates the Python API client for Google Cloud KMS using service account credentials."""
        dictionary_test = json.loads(str(self.service_account))
        if not isinstance(dictionary_test, dict):
            raise Exception("Service Account json is not formatted well. You need to change the json file.")

        credentials_file_name = demisto.uniqueFile() + ".json"
        credentials_file_path = os.path.join(os.getcwd(), credentials_file_name)

        with open(credentials_file_path, "w") as creds_file:
            json_object = json.loads(str(self.service_account))
            json.dump(json_object, creds_file)

        return kms.KeyManagementServiceClient.from_service_account_json(filename=credentials_file_path)  # type: ignore[call-arg]


"""HELPER FUNCTIONS"""


def get_timestamp_seconds(date):
    if not date:
        return 0

    if isinstance(date, timedelta):
        seconds = date.total_seconds()

    else:
        seconds = str(date.timestamp()).split(".")[0]  # type: ignore[assignment]

    return seconds


def get_timestamp_nanoseconds(date):
    if not date:
        return 0

    return str(date.timestamp()).split(".")[1]


def arg_dict_creator(string: Any):
    """Creates a Dict from a CSV string.

    Args:
        string(str): CSV string - formatted as 'field1:value1,field2:value2'.

    Returns:
        Dict from string representation.
    """
    if not string:
        return None

    split_string = string.split(",")
    arg_dict = {}
    for section in split_string:
        section_key, section_value = section.split(":", 1)
        arg_dict[section_key] = section_value

    return arg_dict


def clear_label_commas(labels: Any):
    """When extracted from the response - the labels return with an added commas and spaces
    This function removes these commas.

    Args:
        labels(Dict): a dictionary of labels as returned from the response.

    Returns:
        the label dictionary without the commas and spaces.
    """
    if not labels:
        return None

    cleared_labels = {}  # type:Dict
    for label in labels:
        # A label key can come in the form of: 'info' or _'info' (with an extra space)
        # The following check is whether to drop the first 2 characters or just one
        if str(label).startswith(" "):
            cleared_label_key = label[2:-1]

        else:
            cleared_label_key = label[1:-1]

        if str(labels[label]).startswith(" "):
            cleared_label_value = labels[label][2:-1]

        else:
            cleared_label_value = labels[label][1:-1]

        cleared_labels[cleared_label_key] = cleared_label_value

    return cleared_labels


def key_context_creation(res: Any, project_id: str, location_id: str, key_ring_id: str) -> dict:
    """Creates GoogleKMS.CryptoKey context.

    Args:
        res(Any): `~google.cloud.kms.CryptoKey` instance.
        project_id(str): the project id
        location_id(str): the location id
        key_ring_id(str): the KeyRing id

    Returns:
        Dict representing GoogleKMS.CryptoKey context.
    """
    # remove the CryptoKey path and leave only the name
    pre_name = f"projects/{project_id}/locations/{location_id}/keyRings/{key_ring_id}/cryptoKeys/"
    name = str(res.name).replace(pre_name, "")

    # prepare the labels
    labels = str(res.labels)

    if labels != "{}":
        labels = arg_dict_creator(str(labels)[1:-1])
        labels = clear_label_commas(labels)

    else:
        labels = ""

    key_context = {
        "Name": name,
        "Project": project_id,
        "Location": location_id,
        "KeyRing": key_ring_id,
        "Purpose": kms.CryptoKey.CryptoKeyPurpose(res.purpose).name,
        "CreationTime": datetime.fromtimestamp(int(get_timestamp_seconds(res.create_time))).strftime(DEMISTO_DATETIME_FORMAT),
        "NextRotationTime": datetime.fromtimestamp(int(get_timestamp_seconds(res.next_rotation_time))).strftime(
            DEMISTO_DATETIME_FORMAT
        ),
        "RotationPeriod": f"{get_timestamp_seconds(res.rotation_period)!s}s",
        "Labels": labels,
        "VersionTemplate": {
            "ProtectionLevel": kms.ProtectionLevel(res.version_template.protection_level).name,
            "Algorithm": kms.CryptoKeyVersion.CryptoKeyVersionAlgorithm(res.version_template.algorithm).name,
        },
    }
    # if primary CryptoKeyVersion exists and is returned create context for it.
    # Note: As part of the API - Asymmetric keys do not return primary CryptoKeyVersion info.
    if res.primary and res.primary.name and len(res.primary.name) > 0:
        key_context["PrimaryCryptoKeyVersion"] = {
            "Name": res.primary.name,
            "State": kms.CryptoKeyVersion.CryptoKeyVersionState(res.primary.state).name,
            "CreationTime": datetime.fromtimestamp(int(get_timestamp_seconds(res.primary.create_time))).strftime(
                DEMISTO_DATETIME_FORMAT
            ),
            "ProtectionLevel": kms.ProtectionLevel(res.primary.protection_level).name,
            "Algorithm": kms.CryptoKeyVersion.CryptoKeyVersionAlgorithm(res.primary.algorithm).name,
            "GenerateTime": datetime.fromtimestamp(int(get_timestamp_seconds(res.primary.generate_time))).strftime(
                DEMISTO_DATETIME_FORMAT
            ),
        }

    return key_context


def crypto_key_to_json(crypto_key: Any) -> dict:
    """Creates a json dict from `~google.cloud.kms.CryptoKey` instance to use as raw response.

    Args:
        crypto_key(Any): `~google.cloud.kms.CryptoKey` instance.

    Returns:
        A json Dict containing the raw response.
    """
    # handle labels
    labels = str(crypto_key.labels)
    if labels != "{}":
        labels = arg_dict_creator(str(labels)[1:-1])
        labels = clear_label_commas(labels)

    else:
        labels = ""

    key_json = {
        "name": crypto_key.name,
        "purpose": kms.CryptoKey.CryptoKeyPurpose(crypto_key.purpose).name,
        "create_time": {
            "seconds": get_timestamp_seconds(crypto_key.create_time),
            "nanos": get_timestamp_nanoseconds(crypto_key.create_time),
        },
        "next_rotation_time": {
            "seconds": get_timestamp_seconds(crypto_key.next_rotation_time),
            "nanos": get_timestamp_nanoseconds(crypto_key.next_rotation_time),
        },
        "rotation_period": {"seconds": get_timestamp_seconds(crypto_key.rotation_period)},
        "labels": labels,
        "version_template": {
            "protection_level": kms.ProtectionLevel(crypto_key.version_template.protection_level).name,
            "algorithm": kms.CryptoKeyVersion.CryptoKeyVersionAlgorithm(crypto_key.version_template.algorithm).name,
        },
    }

    if crypto_key.primary:
        key_json["primary"] = {
            "name": crypto_key.primary.name,
            "state": kms.CryptoKeyVersion.CryptoKeyVersionState(crypto_key.primary.state).name,
            "create_time": {
                "seconds": get_timestamp_seconds(crypto_key.primary.create_time),
                "nanos": get_timestamp_nanoseconds(crypto_key.primary.create_time),
            },
            "protection_level": kms.ProtectionLevel(crypto_key.primary.protection_level).name,
            "algorithm": kms.CryptoKeyVersion.CryptoKeyVersionAlgorithm(crypto_key.primary.algorithm).name,
            "generate_time": {
                "seconds": get_timestamp_seconds(crypto_key.primary.generate_time),
                "nanos": get_timestamp_nanoseconds(crypto_key.primary.generate_time),
            },
        }

    return key_json


def demisto_args_extract(client: Client, args: dict[str, Any]) -> tuple[str, str, str, str]:
    """Extracts IDs to use for KMS functions.

    Args:
        args(dict): Demisto arguments.
        client(Client): User Client.

    Returns:
        A tuple containing strings representing the required IDs.
    """
    project_id = client.project

    location_id = args.get("location")
    if location_id == "default":
        location_id = client.location

    key_ring_id = args.get("key_ring")
    if key_ring_id == "default":
        key_ring_id = client.key_ring

    crypto_key_id = args.get("crypto_key")
    return str(project_id), str(location_id), str(key_ring_id), str(crypto_key_id)


def get_update_mask(args: dict[str, Any]) -> dict:
    """Creates the 'updateMask' parameter for the command
    which is a comma separated list of fields to update.

    Args:
        args(dict): Demisto args indicating which field to update.

    Returns:
        A Dict to use as the params for update command.
    """
    update_mask = []
    if args.get("labels") is not None:
        update_mask.append("labels")

    if args.get("next_rotation_time") is not None:
        update_mask.append("next_rotation_time")

    if args.get("purpose") is not None:
        update_mask.append("purpose")

    if args.get("rotation_period") is not None:
        update_mask.append("rotation_period")

    if args.get("attestation") is not None:
        update_mask.append("primary.attestation")

    if args.get("state") is not None:
        update_mask.append("primary.state")

    if args.get("algorithm") is not None:
        update_mask.append("version_template.algorithm")

    if args.get("protection_level") is not None:
        update_mask.append("version_template.protection_level")

    return {"paths": update_mask}


def get_update_command_body(args: dict[str, Any], update_mask: list) -> dict:
    """Creates update command request body, in accordance with the updateMask.

    Args:
        args(Dict): Demisto arguments containing the updated values.
        update_mask(List): List of fields to update.

    Returns:
        Dict to be used as body for update command
    """
    body = {}  # type:Dict[str,Any]
    if "labels" in update_mask:
        # Add label dictionary to body
        body["labels"] = arg_dict_creator(args.get("labels"))

    if "next_rotation_time" in update_mask:
        if str(args.get("next_rotation_time")).isdigit():
            # If next_rotation_time given is a timestamp enter it as is
            body["next_rotation_time"] = {"seconds": int(str(args.get("next_rotation_time")))}

        else:
            # If next_rotation_time is date string, convert it to timestamp
            body["next_rotation_time"] = {
                "seconds": int(datetime.strptime(str(args.get("next_rotation_time")), RFC3339_DATETIME_FORMAT).timestamp())
            }

    if "purpose" in update_mask:
        # Add purpose enum to body
        body["purpose"] = kms.CryptoKey.CryptoKeyPurpose[args["purpose"]].value  # type: ignore[misc]

    if "rotation_period" in update_mask:
        # Add rotation_period to body
        body["rotation_period"] = {"seconds": int(str(args.get("rotation_period")))}

    if "primary.attestation" in update_mask or "primary.state" in update_mask:
        # Init the 'primary' sub-dictionary
        body["primary"] = {}

        if "primary.attestation" in update_mask:
            # Add attestation dict to 'primary' sub-dictionary
            body["primary"]["attestation"] = arg_dict_creator(args.get("attestation"))

        if "primary.state" in update_mask:
            # Add state enum to 'primary' sub-dictionary
            body["primary"]["state"] = kms.CryptoKeyVersion.CryptoKeyVersionState[args.get("state")].value  # type: ignore[misc]

    if "version_template.algorithm" in update_mask or "version_template.protection_level" in update_mask:
        # Init the 'version_template' sun-dictionary
        body["version_template"] = {}

        if "version_template.algorithm" in update_mask:
            # Add algorithm enum to 'version_template' sun-dictionary
            val = kms.CryptoKeyVersion.CryptoKeyVersionAlgorithm[args["algorithm"]].value  # type: ignore[misc]
            body["version_template"]["algorithm"] = val

        if "version_template.protection_level" in update_mask:
            # Add protection_level to 'version_template' sun-dictionary
            val = kms.ProtectionLevel[args["protection_level"]].value  # type: ignore[misc]
            body["version_template"]["protection_level"] = val

    return body


def get_primary_key_version(project_id: str, location_id: str, key_ring_id: str, crypto_key_id: str, client: Client) -> str:
    """Return primary CryptoKeyVersion of a given CryptoKey.

    Args:
        project_id(str): Project of the CryptoKey.
        location_id(str): Location the CryptoKey is assigned to.
        key_ring_id(str): Key Ring in which the CryptoKey exists.
        crypto_key_id(str): The CryptoKey id.
        client(Client): User's Client.

    Returns:
        A string with the full path to the primary CryptoKeyVersion
    """
    # The resource name of the CryptoKey.
    crypto_key_name = client.kms_client.crypto_key_path(project_id, location_id, key_ring_id, crypto_key_id)

    # Get the CryptoKey and extract it's primary version path.
    crypto_key = client.kms_client.get_crypto_key(request={"name": crypto_key_name})
    if crypto_key.primary.name is None:
        raise Exception(f"CryptoKey {crypto_key_name} has no primary CryptoKeyVersion")

    return str(crypto_key.primary.name)


def key_ring_context_and_json_creation(key_ring: Any) -> tuple[dict, dict]:
    key_ring_context = {
        "Name": key_ring.name,
        "CreateTime": datetime.fromtimestamp(int(key_ring.create_time.timestamp())).strftime(DEMISTO_DATETIME_FORMAT),
    }

    key_ring_json = {
        "name": key_ring.name,
        "create_time": {
            "seconds": int(get_timestamp_seconds(key_ring.create_time)),
            "nanos": int(get_timestamp_nanoseconds(key_ring.create_time)),
        },
    }

    return key_ring_context, key_ring_json


"""GENERAL FUNCTIONS"""


def create_crypto_key_command(client: Client, args: dict[str, Any]) -> tuple[str, dict, dict]:
    """Create a new CryptoKey.

    Args:
        client(Client): User Client.
        args(Dict): Demisto arguments.
    """
    project_id, location_id, key_ring_id, crypto_key_id = demisto_args_extract(client, args)

    # The resource name of the KeyRing associated with the CryptoKey.
    key_ring_name = client.kms_client.key_ring_path(project_id, location_id, key_ring_id)

    if args.get("next_rotation_time"):
        # change next_rotation time from date to timestamp if needed.
        if not str(args.get("next_rotation_time")).isdigit():
            next_rotation_time = {
                "seconds": int(datetime.strptime(str(args.get("next_rotation_time")), RFC3339_DATETIME_FORMAT).timestamp())
            }

        else:
            next_rotation_time = {"seconds": int(str(args.get("next_rotation_time")))}

    else:
        # if not next rotation time given - set it to 90 days from now (default by Google)
        next_rotation_time = {"seconds": int((datetime.now() + timedelta(days=90)).timestamp())}

    # Create the CryptoKey object template
    crypto_key = {
        "purpose": kms.CryptoKey.CryptoKeyPurpose[args.get("purpose")].value,  # type: ignore[misc]
        "next_rotation_time": next_rotation_time,
        "labels": arg_dict_creator(args.get("labels")),
        "rotation_period": {"seconds": int(str(args.get("rotation_period")))},
    }

    # Additional info in case CryptoKeyVersion is created
    if args.get("skip_initial_version_creation") != "true":
        crypto_key["primary"] = {
            "state": kms.CryptoKeyVersion.CryptoKeyVersionState[args.get("state")].value,  # type: ignore[misc]
            "attestation": arg_dict_creator(args.get("attestation")),
        }
        crypto_key["version_template"] = {
            "algorithm": kms.CryptoKeyVersion.CryptoKeyVersionAlgorithm[args.get("algorithm")].value,  # type: ignore[misc]
            "protection_level": kms.ProtectionLevel[args.get("protection_level")].value,  # type: ignore[misc]
        }

    # Create a CryptoKey for the given KeyRing.
    response = client.kms_client.create_crypto_key(
        request={
            "parent": key_ring_name,
            "crypto_key_id": crypto_key_id,
            "crypto_key": crypto_key,
            "skip_initial_version_creation": args.get("skip_initial_version_creation") == "true",
        }
    )

    context = key_context_creation(response, project_id, location_id, key_ring_id)

    headers = [
        "CreationTime",
        "Name",
        "Project",
        "Location",
        "KeyRing",
        "Labels",
        "NextRotationTime",
        "Purpose",
        "RotationPeriod",
        "PrimaryCryptoKeyVersion",
        "VersionTemplate",
    ]

    return (
        tableToMarkdown("Google KMS CryptoKey info:", context, removeNull=True, headers=headers),
        {
            f"{INTEGRATION_CONTEXT_NAME}.CryptoKey(val.Name == obj.Name)": context,
        },
        crypto_key_to_json(response),
    )


def symmetric_encrypt_key_command(client: Client, args: dict[str, Any]) -> tuple[str, Any, Any]:
    """Encrypt plaintext to ciphertext using a symmetric key.

    Args:
        client(Client): User Client.
        args(Dict): Demisto arguments.

    Returns:
        The encrypted ciphertext.
    """
    # handle given plaintext - revert it to base 64.
    if args.get("simple_plaintext"):
        plaintext = base64.b64encode(bytes(str(args.get("simple_plaintext")), "utf-8"))

    elif args.get("base64_plaintext"):
        plaintext = base64.b64decode(str(args.get("base64_plaintext")))

    elif args.get("entry_id"):
        file = demisto.getFilePath(args.get("entry_id"))
        file_path = file["path"]
        with open(file_path, "rb") as fp:
            plaintext = base64.b64encode(fp.read())

    else:
        raise ValueError("No object to encrypt.")

    project_id, location_id, key_ring_id, crypto_key_id = demisto_args_extract(client, args)

    additional_authenticated_data = None
    if args.get("additional_authenticated_data"):
        additional_authenticated_data = base64.b64decode(str(args.get("additional_authenticated_data")))

    # The resource name of the CryptoKey.

    crypto_key_name = client.kms_client.crypto_key_path(project_id, location_id, key_ring_id, crypto_key_id)

    # Use the KMS API to encrypt the data.
    response = client.kms_client.encrypt(
        request={"name": crypto_key_name, "plaintext": plaintext, "additional_authenticated_data": additional_authenticated_data}
    )

    # return the created ciphertext cleaned from additional characters.
    ciphertext = str(base64.b64encode(response.ciphertext))[2:-1]

    symmetric_encrypt_context = {
        "CryptoKey": crypto_key_id,
        "IsBase64": args.get("base64_plaintext") is not None,
        "Ciphertext": ciphertext,
    }

    if args.get("entry_id"):
        file_name = demisto.getFilePath(args.get("entry_id"))["name"] + "_encrypted.txt"
        demisto.results(fileResult(file_name, ciphertext))

    return (
        f"The text has been encrypted.\nCiphertext: {ciphertext}",
        {
            f"{INTEGRATION_CONTEXT_NAME}.SymmetricEncrypt(val.CryptoKey == obj.CryptoKey "
            f"&& val.IsBase64 == obj.IsBase64 && val.Ciphertext == obj.Ciphertext)": symmetric_encrypt_context,
        },
        symmetric_encrypt_context,
    )


def symmetric_decrypt_key_command(client: Client, args: dict[str, Any]) -> tuple[str, Any, Any]:
    """Decrypt ciphertext to plaintext using a symmetric key.

    Args:
        client(Client): User Client.
        args(Dict): Demisto agruments.

    Returns:
        The decrypted text.
    """
    if args.get("simple_ciphertext"):
        ciphertext = base64.b64decode(str(args.get("simple_ciphertext")))

    elif args.get("base64_ciphertext"):
        ciphertext = base64.b64decode(str(args.get("base64_ciphertext")))

    elif args.get("entry_id"):
        file = demisto.getFilePath(args.get("entry_id"))
        file_path = file["path"]
        with open(file_path, "rb") as fp:
            ciphertext = base64.b64decode(fp.read())

    else:
        raise ValueError("No object to decrypt.")

    project_id, location_id, key_ring_id, crypto_key_id = demisto_args_extract(client, args)

    additional_authenticated_data = None
    if args.get("additional_authenticated_data"):
        additional_authenticated_data = base64.b64decode(str(args.get("additional_authenticated_data")))

    # The resource name of the CryptoKey.
    crypto_key_name = client.kms_client.crypto_key_path(project_id, location_id, key_ring_id, crypto_key_id)

    # Use the KMS API to decrypt the data.
    response = client.kms_client.decrypt(
        request={
            "name": crypto_key_name,
            "ciphertext": ciphertext,
            "additional_authenticated_data": additional_authenticated_data,
        }
    )

    # handle the resulting plain text if it supposed to be in base64 and clean added characters.
    if args.get("base64_ciphertext"):
        plaintext = str(base64.b64encode(response.plaintext))[2:-1].replace("\\n", "\n")

    elif args.get("simple_ciphertext") or args.get("entry_id"):
        plaintext = str(base64.b64decode(response.plaintext))[2:-1].replace("\\n", "\n")
    else:
        plaintext = ""
        demisto.debug(f"The arguments didn't match any condition, {plaintext=}")

    if args.get("entry_id"):
        file_name = demisto.getFilePath(args.get("entry_id"))["name"] + "_decrypted.txt"
        demisto.results(fileResult(file_name, plaintext))

    symmetric_decrypt_context = {
        "CryptoKey": crypto_key_id,
        "IsBase64": args.get("base64_ciphertext") is not None,
        "Plaintext": plaintext,
    }

    return (
        f"The text has been decrypted.\nPlaintext: {plaintext}",
        {
            f"{INTEGRATION_CONTEXT_NAME}.SymmetricDecrypt(val.CryptoKey == obj.CryptoKey "
            f"&& val.IsBase64 == obj.IsBase64 && val.Plaintext == obj.Plaintext)": symmetric_decrypt_context,
        },
        symmetric_decrypt_context,
    )


def get_key_command(client: Client, args: dict[str, Any]) -> tuple[str, dict, dict]:
    """Gets a CryptoKey.

    Args:
        client(Client): User Client.
        args(Dict): Demisto Arguments.
    """
    project_id, location_id, key_ring_id, crypto_key_id = demisto_args_extract(client, args)

    # The resource name of the CryptoKey.
    crypto_key_name = client.kms_client.crypto_key_path(project_id, location_id, key_ring_id, crypto_key_id)

    # Get CryptoKey info.
    response = client.kms_client.get_crypto_key(request={"name": crypto_key_name})

    context = key_context_creation(response, project_id, location_id, key_ring_id)

    headers = [
        "CreationTime",
        "Name",
        "Project",
        "Location",
        "KeyRing",
        "Labels",
        "NextRotationTime",
        "Purpose",
        "RotationPeriod",
        "PrimaryCryptoKeyVersion",
        "VersionTemplate",
    ]

    return (
        tableToMarkdown("Google KMS CryptoKey info:", context, removeNull=True, headers=headers),
        {
            f"{INTEGRATION_CONTEXT_NAME}.CryptoKey(val.Name == obj.Name)": context,
        },
        crypto_key_to_json(response),
    )


def disable_key_command(client: Client, args: dict[str, Any]) -> tuple[str, Any, Any]:
    """Disable a given CryptoKeyVersion.

    Args:
        client(Client): User Client.
        args(Dict): Demisto arguments.
    """
    project_id, location_id, key_ring_id, crypto_key_id = demisto_args_extract(client, args)
    crypto_key_version = args.get("crypto_key_version")

    if crypto_key_version == "default":
        # if no CryptoKeyVersion given extract the primary CryptoKeyVersion.
        crypto_key_version_name = get_primary_key_version(project_id, location_id, key_ring_id, crypto_key_id, client)

    else:
        # Construct the resource name of the CryptoKeyVersion.
        crypto_key_version_name = client.kms_client.crypto_key_version_path(
            project_id, location_id, key_ring_id, crypto_key_id, crypto_key_version
        )

    # if not CryptoKeyVersion is given nor was is extracted from the CryptoKey - raise error to the user.
    if crypto_key_version_name is None or len(crypto_key_version_name) == 0:
        raise Exception("Please insert primary CryptoKeyVersion ID")

    # Use the KMS API to disable the CryptoKeyVersion.
    new_state = kms.CryptoKeyVersion.CryptoKeyVersionState.DISABLED
    version = {"name": crypto_key_version_name, "state": new_state}
    update_mask = {"paths": ["state"]}

    # Print results
    response = client.kms_client.update_crypto_key_version(request={"crypto_key_version": version, "update_mask": update_mask})
    return (
        f"CryptoKeyVersion {crypto_key_version_name}'s state has been set to "
        f"{kms.CryptoKeyVersion.CryptoKeyVersionState(response.state).name}.",
        None,
        None,
    )


def enable_key_command(client: Client, args: dict[str, Any]) -> tuple[str, Any, Any]:
    """Enable a CryptoKeyVersion.

    Args:
        client(Client): User Client.
        args(Dict): Demisto arguments.
    """
    project_id, location_id, key_ring_id, crypto_key_id = demisto_args_extract(client, args)
    crypto_key_version = args.get("crypto_key_version")

    if crypto_key_version == "default":
        # if no CryptoKeyVersion given extract the primary CryptoKeyVersion.
        crypto_key_version_name = get_primary_key_version(project_id, location_id, key_ring_id, crypto_key_id, client)

    else:
        # Construct the resource name of the CryptoKeyVersion.
        crypto_key_version_name = client.kms_client.crypto_key_version_path(
            project_id, location_id, key_ring_id, crypto_key_id, crypto_key_version
        )

    # if not CryptoKeyVersion is given nor was is extracted from the CryptoKey - raise error to the user.
    if crypto_key_version_name is None or len(crypto_key_version_name) == 0:
        raise Exception("Please insert primary CryptoKeyVersion ID")

    # Use the KMS API to enable the CryptoKeyVersion.
    new_state = kms.CryptoKeyVersion.CryptoKeyVersionState.ENABLED
    version = {"name": crypto_key_version_name, "state": new_state}
    update_mask = {"paths": ["state"]}

    # Print results
    response = client.kms_client.update_crypto_key_version(request={"crypto_key_version": version, "update_mask": update_mask})
    return (
        f"CryptoKeyVersion {crypto_key_version_name}'s state has been set to "
        f"{kms.CryptoKeyVersion.CryptoKeyVersionState(response.state).name}.",
        None,
        None,
    )


def destroy_key_command(client: Client, args: dict[str, Any]) -> tuple[str, Any, Any]:
    """Schedule the destruction of a given CryptoKeyVersion.

    Args:
        client(Client): User Client.
        args(Dict): Demisto arguments.
    """
    project_id, location_id, key_ring_id, crypto_key_id = demisto_args_extract(client, args)
    crypto_key_version = args.get("crypto_key_version")

    if crypto_key_version == "default":
        # if no CryptoKeyVersion given extract the primary CryptoKeyVersion.
        crypto_key_version_name = get_primary_key_version(project_id, location_id, key_ring_id, crypto_key_id, client)

    else:
        # Construct the resource name of the CryptoKeyVersion.
        crypto_key_version_name = client.kms_client.crypto_key_version_path(
            project_id, location_id, key_ring_id, crypto_key_id, crypto_key_version
        )

    # if not CryptoKeyVersion is given nor was is extracted from the CryptoKey - raise error to the user.
    if crypto_key_version_name is None or len(crypto_key_version_name) == 0:
        raise Exception("Please insert primary CryptoKeyVersion ID")

    # Use the KMS API to mark the CryptoKeyVersion for destruction.
    response = client.kms_client.destroy_crypto_key_version(request={"name": crypto_key_version_name})

    # Print results
    return (
        f"CryptoKeyVersion {crypto_key_version_name}'s state has been set to "
        f"{kms.CryptoKeyVersion.CryptoKeyVersionState(response.state).name}, it will be destroyed in 24h.",
        None,
        None,
    )


def restore_key_command(client: Client, args: dict[str, Any]) -> tuple[str, Any, Any]:
    """Restores a CryptoKeyVersion scheduled for destruction.

    Args:
        client(Client): User Client.
        args(Dict): Demisto arguments.
    """
    project_id, location_id, key_ring_id, crypto_key_id = demisto_args_extract(client, args)
    crypto_key_version = args.get("crypto_key_version")

    if crypto_key_version == "default":
        # if no CryptoKeyVersion given extract the primary CryptoKeyVersion.
        crypto_key_version_name = get_primary_key_version(project_id, location_id, key_ring_id, crypto_key_id, client)

    else:
        # Construct the resource name of the CryptoKeyVersion.
        crypto_key_version_name = client.kms_client.crypto_key_version_path(
            project_id, location_id, key_ring_id, crypto_key_id, crypto_key_version
        )

    # if not CryptoKeyVersion is given nor was is extracted from the CryptoKey - raise error to the user.
    if crypto_key_version_name is None or len(crypto_key_version_name) == 0:
        raise Exception("Please insert primary CryptoKeyVersion ID")

    # Use the KMS API to restore the CryptoKeyVersion.
    response = client.kms_client.restore_crypto_key_version(request={"name": crypto_key_version_name})

    # Print results
    return (
        f"CryptoKeyVersion {crypto_key_version_name}'s state has been set to "
        f"{kms.CryptoKeyVersion.CryptoKeyVersionState(response.state).name}.",
        None,
        None,
    )


def update_key_command(client: Client, args: dict[str, Any]) -> tuple[str, dict, dict]:
    """Update a given CryptoKey.

    Args:
        client(Client): User Client.
        args(Dict): Demisto arguments.
    """
    project_id, location_id, key_ring_id, crypto_key_id = demisto_args_extract(client, args)

    # The resource name of the CryptoKey.
    crypto_key_name = client.kms_client.crypto_key_path(project_id, location_id, key_ring_id, crypto_key_id)

    # create a list of fields to be updated using the command - the field is called update mask in the API.
    update_mask = get_update_mask(args)

    # create the body of the update request.
    crypto_key = get_update_command_body(args=args, update_mask=update_mask["paths"])
    crypto_key["name"] = crypto_key_name

    # update command using the KMS API.
    response = client.kms_client.update_crypto_key(request={"crypto_key": crypto_key, "update_mask": update_mask})

    context = key_context_creation(response, project_id, location_id, key_ring_id)

    headers = [
        "CreationTime",
        "Name",
        "Project",
        "Location",
        "KeyRing",
        "Labels",
        "NextRotationTime",
        "Purpose",
        "RotationPeriod",
        "PrimaryCryptoKeyVersion",
        "VersionTemplate",
    ]

    return (
        tableToMarkdown("Google KMS CryptoKey info:", context, removeNull=True, headers=headers),
        {
            f"{INTEGRATION_CONTEXT_NAME}.CryptoKey(val.Name == obj.Name)": context,
        },
        crypto_key_to_json(response),
    )


def list_keys_command(client: Client, args: dict[str, Any]) -> tuple[str, dict, list]:
    """List All keys in a KeyRing.

    Args:
        client(Client): User Client.
        args(Dict): Demisto arguments.

    """
    project_id, location_id, key_ring_id, _ = demisto_args_extract(client, args)

    # Get the full path to the KeyRing in which to list the keys.
    key_ring_name = client.kms_client.key_ring_path(project_id, location_id, key_ring_id)

    # if needed add state filter.
    filter_state = args.get("key_state", None)

    response = client.kms_client.list_crypto_keys(request={"parent": key_ring_name, "filter": filter_state})

    overall_context = []  # type: List
    overall_raw = []  # type: List
    for crypto_key in response:
        overall_context.append(key_context_creation(crypto_key, project_id, location_id, key_ring_id))
        overall_raw.append(crypto_key_to_json(crypto_key))

    headers = [
        "CreationTime",
        "Name",
        "Project",
        "Location",
        "KeyRing",
        "Labels",
        "NextRotationTime",
        "Purpose",
        "RotationPeriod",
        "PrimaryCryptoKeyVersion",
        "VersionTemplate",
    ]

    return (
        tableToMarkdown(name="CryptoKeys:", t=overall_context, removeNull=True, headers=headers),
        {
            f"{INTEGRATION_CONTEXT_NAME}.CryptoKey(val.Name == obj.Name)": overall_context,
        },
        overall_raw,
    )


def asymmetric_encrypt_command(client: Client, args: dict[str, Any]) -> tuple[str, Any, Any]:
    """Encrypt plainttext using an asymmetric key.

    Args:
        client(Client): User's client.
        args(dict): Demisto arguments.

    Returns:
        The encrypted ciphertext.
    """
    # handle the plaintext - convert to base64
    if args.get("simple_plaintext"):
        plaintext = base64.b64encode(bytes(str(args.get("simple_plaintext")), "utf-8"))

    elif args.get("base64_plaintext"):
        plaintext = base64.b64decode(str(args.get("base64_plaintext")))

    elif args.get("entry_id"):
        file = demisto.getFilePath(args.get("entry_id"))
        file_path = file["path"]
        with open(file_path, "rb") as fp:
            plaintext = base64.b64encode(fp.read())

    else:
        raise ValueError("No object to encrypt.")

    project_id, location_id, key_ring_id, crypto_key_id = demisto_args_extract(client, args)
    crypto_key_version = args.get("crypto_key_version")

    # Construct the resource name of the CryptoKeyVersion.
    crypto_key_version_name = client.kms_client.crypto_key_version_path(
        project_id, location_id, key_ring_id, crypto_key_id, crypto_key_version
    )
    # get the CryptoKeyVersion info and check it's algorithm.
    crypto_key_version_info = client.kms_client.get_crypto_key_version(request={"name": crypto_key_version_name})
    key_algo = kms.CryptoKeyVersion.CryptoKeyVersionAlgorithm(crypto_key_version_info.algorithm).name

    # Algorithm must be a "DECRYPT" type asymmetric algorithm - if not, raise an error to the user.
    if "DECRYPT" not in key_algo:
        raise ValueError(f"{crypto_key_version_name} is not a valid asymmetric CryptoKeyVersion")

    # get public key of the asymmetric encryption.
    public_key_response = client.kms_client.get_public_key(request={"name": crypto_key_version_name})
    key_txt = public_key_response.pem.encode("ascii")
    public_key = serialization.load_pem_public_key(key_txt, default_backend())

    # using the CryptoKeyVersion algorithm - create the necessary padding for the encryption.
    if "SHA256" in key_algo:
        # create padding with SHA256
        pad = padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
    else:
        # create padding with SHA512
        pad = padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA512()), algorithm=hashes.SHA512(), label=None)

    # encrypt plaintext and return the cipertext without added characters.
    ciphertext = str(base64.b64encode(public_key.encrypt(plaintext, pad)))[2:-1]  # type: ignore[union-attr]

    asymmetric_encrypt_context = {
        "CryptoKey": crypto_key_id,
        "IsBase64": args.get("base64_plaintext") is not None,
        "Ciphertext": ciphertext,
    }

    if args.get("entry_id"):
        file_name = demisto.getFilePath(args.get("entry_id"))["name"] + "_encrypted.txt"
        demisto.results(fileResult(file_name, ciphertext))

    return (
        f"The text has been encrypted.\nCiphertext: {ciphertext}",
        {
            f"{INTEGRATION_CONTEXT_NAME}.AsymmetricEncrypt(val.CryptoKey == obj.CryptoKey "
            f"&& val.IsBase64 == obj.IsBase64 && val.Ciphertext == obj.Ciphertext)": asymmetric_encrypt_context,
        },
        asymmetric_encrypt_context,
    )


def asymmetric_decrypt_command(client: Client, args: dict[str, Any]) -> tuple[str, Any, Any]:
    """Decrypt chipertext to plaintext using asymmetric key.

    Args:
        client(Client): User's Client.
        args(dict): Demisto arguments.

    Returns:
        The decrypted plaintext.
    """
    project_id, location_id, key_ring_id, crypto_key_id = demisto_args_extract(client, args)
    crypto_key_version = args.get("crypto_key_version")

    # Construct the resource name of the CryptoKeyVersion.
    crypto_key_version_name = client.kms_client.crypto_key_version_path(
        project_id, location_id, key_ring_id, crypto_key_id, crypto_key_version
    )

    if args.get("simple_ciphertext"):
        ciphertext = base64.b64decode(str(args.get("simple_ciphertext")))

    elif args.get("base64_ciphertext"):
        ciphertext = base64.b64decode(str(args.get("base64_ciphertext")))

    elif args.get("entry_id"):
        file = demisto.getFilePath(args.get("entry_id"))
        file_path = file["path"]
        with open(file_path, "rb") as fp:
            ciphertext = base64.b64decode(fp.read())

    else:
        raise ValueError("No object to decrypt.")

    response = client.kms_client.asymmetric_decrypt(request={"name": crypto_key_version_name, "ciphertext": ciphertext})

    # handle the created plaintext back to base64 if needed and clear added characters.
    if args.get("base64_ciphertext"):
        plaintext = str(base64.b64encode(response.plaintext))[2:-1].replace("\\n", "\n")

    elif args.get("simple_ciphertext") or args.get("entry_id"):
        plaintext = str(base64.b64decode(response.plaintext))[2:-1].replace("\\n", "\n")

    else:
        plaintext = ""
        demisto.debug(f"The arguments didn't match any condition, {plaintext=}")

    if args.get("entry_id"):
        file_name = demisto.getFilePath(args.get("entry_id"))["name"] + "_decrypted.txt"
        demisto.results(fileResult(file_name, plaintext))

    asymmetric_decrypt_context = {
        "CryptoKey": crypto_key_id,
        "IsBase64": args.get("base64_ciphertext") is not None,
        "Plaintext": plaintext,
    }

    return (
        f"The text has been decrypted.\nPlaintext: {plaintext}",
        {
            f"{INTEGRATION_CONTEXT_NAME}.AsymmetricDecrypt(val.CryptoKey == obj.CryptoKey "
            f"&& val.IsBase64 == obj.IsBase64 && val.Plaintext == obj.Plaintext)": asymmetric_decrypt_context,
        },
        asymmetric_decrypt_context,
    )


def list_key_rings_command(client: Client, args: dict[str, Any]) -> tuple[str, Any, Any]:
    """List all KeyRings in a given location

    Args:
        client(Client): User's client.
        args(dict): Demisto args.
    """
    # listing the KeyRings in order to check responses from the API.
    locations = []
    if args.get("location") == "default":
        locations.append(client.location)

    else:
        locations.append(args.get("location"))

    # paramater 'all' checks all possible locations
    if args.get("all") == "yes":
        locations = [
            "global",
            "asia-east1",
            "asia-east2",
            "asia-northeast1",
            "asia-northeast2",
            "asia-south1",
            "asia-southeast1",
            "australia-southeast1",
            "europe-north1",
            "europe-west1",
            "europe-west2",
            "europe-west3",
            "europe-west4",
            "europe-west6",
            "northamerica-northeast1",
            "us-central1",
            "us-east1",
            "us-east4",
            "us-west1",
            "us-west2",
            "southamerica-east1",
            "eur4",
            "nam4",
            "asia",
            "europe",
            "us",
        ]

    key_rings_context = []
    key_rings_json = []

    for location in locations:
        location_path = f"projects/{client.project}/locations/{location}"
        # the response is a iterable containing the KeyRings info.
        response = client.kms_client.list_key_rings(request={"parent": location_path})

        for key_ring in list(response):
            single_context, single_json = key_ring_context_and_json_creation(key_ring)
            key_rings_context.append(single_context)
            key_rings_json.append(single_json)

    return (
        tableToMarkdown(name="KeyRings:", t=key_rings_context, removeNull=True),
        {f"{INTEGRATION_CONTEXT_NAME}.KeyRing(val.Name == obj.Name)": key_rings_context},
        key_rings_json,
    )


def list_all_keys_command(client: Client, args: dict[str, Any]) -> tuple[str, Any, Any]:
    """List all CryptokKeys across all KeyRings in a given location.

    Args:
        client(Client): User's client.
        args(dict): Demisto args.
    """
    locations = []
    if args.get("location") == "default":
        locations.append(client.location)

    else:
        locations.append(args.get("location"))

    # paramater 'all' checks all possible locations
    if args.get("all") == "yes":
        locations = [
            "global",
            "asia-east1",
            "asia-east2",
            "asia-northeast1",
            "asia-northeast2",
            "asia-south1",
            "asia-southeast1",
            "australia-southeast1",
            "europe-north1",
            "europe-west1",
            "europe-west2",
            "europe-west3",
            "europe-west4",
            "europe-west6",
            "northamerica-northeast1",
            "us-central1",
            "us-east1",
            "us-east4",
            "us-west1",
            "us-west2",
            "southamerica-east1",
            "eur4",
            "nam4",
            "asia",
            "europe",
            "us",
        ]

    keys_context = []
    keys_json = []
    filter_state = args.get("key_state")
    for location in locations:
        location_path = f"projects/{client.project}/locations/{location}"
        # the response is a iterable containing the KeyRings info.
        response = client.kms_client.list_key_rings(request={"parent": location_path})

        for key_ring in list(response):
            key_ring_name = key_ring.name
            pre_name = f"projects/{client.project}/locations/{location}/keyRings/"
            key_ring_id = str(key_ring_name).replace(pre_name, "")
            crypto_key_response = client.kms_client.list_crypto_keys(request={"parent": key_ring_name, "filter": filter_state})

            for crypto_key in crypto_key_response:
                keys_context.append(key_context_creation(crypto_key, str(client.project), str(location), str(key_ring_id)))
                keys_json.append(crypto_key_to_json(crypto_key))

    headers = [
        "CreationTime",
        "Name",
        "Project",
        "Location",
        "KeyRing",
        "Labels",
        "NextRotationTime",
        "Purpose",
        "RotationPeriod",
        "PrimaryCryptoKeyVersion",
        "VersionTemplate",
    ]

    return (
        tableToMarkdown(name="CryptoKeys:", t=keys_context, removeNull=True, headers=headers),
        {
            f"{INTEGRATION_CONTEXT_NAME}.CryptoKey(val.Name == obj.Name && val.Location == obj.Location"
            f"&& val.KeyRing == obj. KeyRing)": keys_context,
        },
        keys_json,
    )


def get_public_key_command(client: Client, args: dict[str, Any]) -> tuple[str, dict, dict]:
    """Get the public key from an asymmetric CryptoKey

    Args:
        client(Client): User's client.
        args(dict): Demisto args.
    """
    project_id, location_id, key_ring_id, crypto_key_id = demisto_args_extract(client, args)
    crypto_key_version = args.get("crypto_key_version")

    # Construct the resource name of the CryptoKeyVersion.
    crypto_key_version_name = client.kms_client.crypto_key_version_path(
        project_id, location_id, key_ring_id, crypto_key_id, crypto_key_version
    )

    public_key_response = client.kms_client.get_public_key(request={"name": crypto_key_version_name})

    public_key_context = {
        "CryptoKey": crypto_key_id,
        "PEM": str(public_key_response.pem),
        "Algorithm": kms.CryptoKeyVersion.CryptoKeyVersionAlgorithm(public_key_response.algorithm).name,
    }

    return (
        f"The Public Key for CryptoKey {crypto_key_id} is:\n{public_key_context.get('PEM')}",
        {f"{INTEGRATION_CONTEXT_NAME}.PublicKey(val.CryptoKey == obj.CryptoKey && val.PEM == obj.PEM": public_key_context},
        public_key_context,
    )


def test_function(client: Client) -> None:
    """Test's user's input this checks if the given service account has any of the required permissions
    to use the integration

    In client creation we check that the entered service account is a valid json and has all the required
    fields to create the client.

    In this function we try and get a response from Google KMS just to make sure we can connect to it.

    This test does NOT check if the service account has the required permissions to get a VALID response from KMS.

    Args:
        client(Client): User Client.
    """
    # creating a valid resource name
    if client.key_ring:
        key_ring = client.key_ring

    else:
        key_ring = "random"

    key_ring_name = client.kms_client.key_ring_path(client.project, client.location, key_ring)
    client.kms_client.list_crypto_keys(request={"parent": key_ring_name})


def main():
    COMMANDS = {
        f"{INTEGRATION_COMMAND_NAME}-create-key": create_crypto_key_command,
        f"{INTEGRATION_COMMAND_NAME}-symmetric-decrypt": symmetric_decrypt_key_command,
        f"{INTEGRATION_COMMAND_NAME}-symmetric-encrypt": symmetric_encrypt_key_command,
        f"{INTEGRATION_COMMAND_NAME}-get-key": get_key_command,
        f"{INTEGRATION_COMMAND_NAME}-update-key": update_key_command,
        f"{INTEGRATION_COMMAND_NAME}-destroy-key": destroy_key_command,
        f"{INTEGRATION_COMMAND_NAME}-restore-key": restore_key_command,
        f"{INTEGRATION_COMMAND_NAME}-disable-key": disable_key_command,
        f"{INTEGRATION_COMMAND_NAME}-enable-key": enable_key_command,
        f"{INTEGRATION_COMMAND_NAME}-list-keys": list_keys_command,
        f"{INTEGRATION_COMMAND_NAME}-asymmetric-encrypt": asymmetric_encrypt_command,
        f"{INTEGRATION_COMMAND_NAME}-asymmetric-decrypt": asymmetric_decrypt_command,
        f"{INTEGRATION_COMMAND_NAME}-list-key-rings": list_key_rings_command,
        f"{INTEGRATION_COMMAND_NAME}-list-all-keys": list_all_keys_command,
        f"{INTEGRATION_COMMAND_NAME}-get-public-key": get_public_key_command,
    }

    command = demisto.command()
    LOG(f"{INTEGRATION_NAME}: command is {command}")
    try:
        client = Client(demisto.params())

        if command == "test-module":
            test_function(client)
            demisto.results("ok")

        if command not in COMMANDS:
            raise NotImplementedError(f'Command "{command}" is not implemented.')

        cmd_func = COMMANDS.get(command)  # type: ignore

        results = cmd_func(client, demisto.args())  # type: ignore

        return_outputs(*results)

    except Exception as e:
        return_error(f"{INTEGRATION_NAME}: {e!s}", e)


if __name__ in ["__main__", "builtin", "builtins"]:
    main()