privateca/snippets/conftest.py

# Copyright 2021 Google LLC
#
# 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.

from multiprocessing import Pool
import os
import random
import uuid

import backoff
from google.api_core.exceptions import FailedPrecondition
from google.api_core.exceptions import ServiceUnavailable
import google.auth
from google.cloud.security import privateca_v1
import pytest

from create_ca_pool import create_ca_pool
from create_certificate_authority import create_certificate_authority
from create_certificate_template import create_certificate_template
from delete_certificate_authority import delete_certificate_authority
from delete_certificate_template import delete_certificate_template
from disable_certificate_authority import disable_certificate_authority
from enable_certificate_authority import enable_certificate_authority

PROJECT = google.auth.default()[1]
LOCATIONS = (
    "us-central1",
    "europe-north1",
    "europe-central2",
    "europe-west2",
    "us-east4",
    "europe-west1",
)
LOCATION = random.choice(LOCATIONS)
COMMON_NAME = "COMMON_NAME"
ORGANIZATION = "ORGANIZATION"
CA_DURATION = 1000000


def delete_cas_from_pool(ca_pool_name: str) -> None:
    client = privateca_v1.CertificateAuthorityServiceClient()
    for ca in client.list_certificate_authorities(parent=ca_pool_name):
        # Check if the CA is enabled.
        if ca.state == privateca_v1.CertificateAuthority.State.ENABLED:
            request = privateca_v1.DisableCertificateAuthorityRequest(name=ca.name)
            client.disable_certificate_authority(request=request)

        # Delete CA.
        ca_state = client.get_certificate_authority(name=ca.name).state
        try:
            if ca_state != privateca_v1.CertificateAuthority.State.DELETED:
                delete_ca_request = privateca_v1.DeleteCertificateAuthorityRequest()
                delete_ca_request.name = ca.name
                delete_ca_request.ignore_active_certificates = True
                delete_ca_request.skip_grace_period = True
                client.delete_certificate_authority(request=delete_ca_request).result(
                    timeout=300
                )
                print(f" * {ca.name} Deleted!")
        except FailedPrecondition as err:
            print(err)
            continue


def delete_one_pool(ca_pool_name: str) -> None:
    client = privateca_v1.CertificateAuthorityServiceClient()
    try:
        delete_ca_pool_request = privateca_v1.DeleteCaPoolRequest()
        delete_ca_pool_request.name = ca_pool_name
        print(f"Deleting {ca_pool_name}")
        client.delete_ca_pool(request=delete_ca_pool_request).result(timeout=300)
    except FailedPrecondition:
        print(f"Precondition failed for {ca_pool_name} :(")


@backoff.on_exception(backoff.expo, ServiceUnavailable, max_tries=3)
def delete_stale_resources() -> None:
    client = privateca_v1.CertificateAuthorityServiceClient()
    pool_names = []
    for location in LOCATIONS:
        location_path = client.common_location_path(PROJECT, location)
        request = privateca_v1.ListCaPoolsRequest(parent=location_path)
        for ca_pool in client.list_ca_pools(request=request):
            pool_names.append(ca_pool.name)

    with Pool(max(2, os.cpu_count() - 2)) as p:
        print(f"Going to clean up CAs from {len(pool_names)} pools.")
        p.map(delete_cas_from_pool, pool_names)

    with Pool(max(2, os.cpu_count() - 2)) as p:
        print(f"Going to delete {len(pool_names)} pools.")
        p.map(delete_one_pool, pool_names)


def generate_name() -> str:
    return "test-" + uuid.uuid4().hex[:10]


@pytest.fixture
def ca_pool(ca_pool_autodelete_name):
    create_ca_pool(PROJECT, LOCATION, ca_pool_autodelete_name)
    yield ca_pool_autodelete_name


@pytest.fixture
def certificate_authority(ca_pool):
    CA_NAME = generate_name()

    create_certificate_authority(
        PROJECT, LOCATION, ca_pool, CA_NAME, COMMON_NAME, ORGANIZATION, CA_DURATION
    )

    yield ca_pool, CA_NAME

    # CA Pool cleanup will remove the certificate.
    # delete_certificate_authority(PROJECT, LOCATION, ca_pool, CA_NAME)


@pytest.fixture
def deleted_certificate_authority(ca_pool):
    CA_NAME = generate_name()

    create_certificate_authority(
        PROJECT, LOCATION, ca_pool, CA_NAME, COMMON_NAME, ORGANIZATION, CA_DURATION
    )

    enable_certificate_authority(PROJECT, LOCATION, ca_pool, CA_NAME)
    disable_certificate_authority(PROJECT, LOCATION, ca_pool, CA_NAME)
    delete_certificate_authority(PROJECT, LOCATION, ca_pool, CA_NAME)

    yield ca_pool, CA_NAME


@pytest.fixture
def certificate_template():
    TEMPLATE_NAME = generate_name()

    create_certificate_template(PROJECT, LOCATION, TEMPLATE_NAME)

    yield TEMPLATE_NAME

    delete_certificate_template(PROJECT, LOCATION, TEMPLATE_NAME)


@pytest.fixture
def ca_pool_autodelete_name():
    name = generate_name()
    yield name
    ca_client = privateca_v1.CertificateAuthorityServiceClient()

    ca_pool_path = ca_client.ca_pool_path(PROJECT, LOCATION, name)
    delete_cas_from_pool(ca_pool_path)
    delete_one_pool(ca_pool_path)


@pytest.fixture
def ca_pool_autodelete_name2():
    name = generate_name()
    yield name
    ca_client = privateca_v1.CertificateAuthorityServiceClient()

    ca_pool_path = ca_client.ca_pool_path(PROJECT, LOCATION, name)
    delete_cas_from_pool(ca_pool_path)
    delete_one_pool(ca_pool_path)