diff --git a/functions/billing_stop_on_notification/requirements-test.txt b/functions/billing_stop_on_notification/requirements-test.txt new file mode 100644 index 0000000000..66801836e2 --- /dev/null +++ b/functions/billing_stop_on_notification/requirements-test.txt @@ -0,0 +1,2 @@ +pytest==8.3.5 +cloudevents==1.11.0 \ No newline at end of file diff --git a/functions/billing_stop_on_notification/requirements.txt b/functions/billing_stop_on_notification/requirements.txt new file mode 100644 index 0000000000..912b07cd0a --- /dev/null +++ b/functions/billing_stop_on_notification/requirements.txt @@ -0,0 +1,3 @@ +functions-framework==3.* +google-cloud-billing==1.16.2 +google-cloud-logging==3.12.1 diff --git a/functions/billing_stop_on_notification/stop_billing.py b/functions/billing_stop_on_notification/stop_billing.py new file mode 100644 index 0000000000..fcb6563e05 --- /dev/null +++ b/functions/billing_stop_on_notification/stop_billing.py @@ -0,0 +1,169 @@ +# Copyright 2025 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. + +# [START functions_billing_stop] +# WARNING: The following action, if not in simulation mode, will disable billing +# for the project, potentially stopping all services and causing outages. +# Ensure thorough testing and understanding before enabling live deactivation. + +import base64 +import json +import os +import urllib.request + +from cloudevents.http.event import CloudEvent +import functions_framework + +from google.api_core import exceptions +from google.cloud import billing_v1 +from google.cloud import logging + +billing_client = billing_v1.CloudBillingClient() + + +def get_project_id() -> str: + """Retrieves the Google Cloud Project ID. + + This function first attempts to get the project ID from the + `GOOGLE_CLOUD_PROJECT` environment variable. If the environment + variable is not set or is None, it then attempts to retrieve the + project ID from the Google Cloud metadata server. + + Returns: + str: The Google Cloud Project ID. + + Raises: + ValueError: If the project ID cannot be determined either from + the environment variable or the metadata server. + """ + + # Read the environment variable, usually set manually + project_id = os.getenv("GOOGLE_CLOUD_PROJECT") + if project_id is not None: + return project_id + + # Otherwise, get the `project-id`` from the Metadata server + url = "http://metadata.google.internal/computeMetadata/v1/project/project-id" + req = urllib.request.Request(url) + req.add_header("Metadata-Flavor", "Google") + project_id = urllib.request.urlopen(req).read().decode() + + if project_id is None: + raise ValueError("project-id metadata not found.") + + return project_id + + +@functions_framework.cloud_event +def stop_billing(cloud_event: CloudEvent) -> None: + # TODO(developer): As stoping billing is a destructive action + # for your project, change the following constant to False + # after you validate with a test budget. + SIMULATE_DEACTIVATION = True + + PROJECT_ID = get_project_id() + PROJECT_NAME = f"projects/{PROJECT_ID}" + + event_data = base64.b64decode( + cloud_event.data["message"]["data"] + ).decode("utf-8") + + event_dict = json.loads(event_data) + cost_amount = event_dict["costAmount"] + budget_amount = event_dict["budgetAmount"] + print(f"Cost: {cost_amount} Budget: {budget_amount}") + + if cost_amount <= budget_amount: + print("No action required. Current cost is within budget.") + return + + print(f"Disabling billing for project '{PROJECT_NAME}'...") + + is_billing_enabled = _is_billing_enabled(PROJECT_NAME) + + if is_billing_enabled: + _disable_billing_for_project( + PROJECT_NAME, + SIMULATE_DEACTIVATION + ) + else: + print("Billing is already disabled.") + + +def _is_billing_enabled(project_name: str) -> bool: + """Determine whether billing is enabled for a project. + + Args: + project_name: Name of project to check if billing is enabled. + + Returns: + Whether project has billing enabled or not. + """ + try: + print(f"Getting billing info for project '{project_name}'...") + response = billing_client.get_project_billing_info(name=project_name) + + return response.billing_enabled + except Exception as e: + print(f'Error getting billing info: {e}') + print( + "Unable to determine if billing is enabled on specified project, " + "assuming billing is enabled." + ) + + return True + + +def _disable_billing_for_project( + project_name: str, + simulate_deactivation: bool, +) -> None: + """Disable billing for a project by removing its billing account. + + Args: + project_name: Name of project to disable billing. + simulate_deactivation: + If True, it won't actually disable billing. + Useful to validate with test budgets. + """ + + # Log this operation in Cloud Logging + logging_client = logging.Client() + logger = logging_client.logger(name="disable-billing") + + if simulate_deactivation: + entry_text = "Billing disabled. (Simulated)" + print(entry_text) + logger.log_text(entry_text, severity="CRITICAL") + return + + # Find more information about `updateBillingInfo` API method here: + # https://cloud.google.com/billing/docs/reference/rest/v1/projects/updateBillingInfo + try: + # To disable billing set the `billing_account_name` field to empty + project_billing_info = billing_v1.ProjectBillingInfo( + billing_account_name="" + ) + + response = billing_client.update_project_billing_info( + name=project_name, + project_billing_info=project_billing_info + ) + + entry_text = f"Billing disabled: {response}" + print(entry_text) + logger.log_text(entry_text, severity="CRITICAL") + except exceptions.PermissionDenied: + print("Failed to disable billing, check permissions.") +# [END functions_billing_stop] diff --git a/functions/billing_stop_on_notification/stop_billing_test.py b/functions/billing_stop_on_notification/stop_billing_test.py new file mode 100644 index 0000000000..5ad4f9f3bf --- /dev/null +++ b/functions/billing_stop_on_notification/stop_billing_test.py @@ -0,0 +1,83 @@ +# Copyright 2025 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. + +import base64 +import json + +from cloudevents.conversion import to_structured +from cloudevents.http import CloudEvent + +from flask.testing import FlaskClient + +from functions_framework import create_app + +import pytest + + +@pytest.fixture +def cloud_event_budget_alert() -> CloudEvent: + attributes = { + "specversion": "1.0", + "id": "my-id", + "source": "//pubsub.googleapis.com/projects/PROJECT_NAME/topics/TOPIC_NAME", + "type": "google.cloud.pubsub.topic.v1.messagePublished", + "datacontenttype": "application/json", + "time": "2025-05-09T18:32:46.572Z" + } + + budget_data = { + "budgetDisplayName": "BUDGET_NAME", + "alertThresholdExceeded": 1.0, + "costAmount": 2.0, + "costIntervalStart": "2025-05-01T07:00:00Z", + "budgetAmount": 0.01, + "budgetAmountType": "SPECIFIED_AMOUNT", + "currencyCode": "USD" + } + + json_string = json.dumps(budget_data) + message_base64 = base64.b64encode(json_string.encode('utf-8')).decode('utf-8') + + data = { + "message": { + "data": message_base64 + } + } + + return CloudEvent(attributes, data) + + +@pytest.fixture +def client() -> FlaskClient: + source = "stop_billing.py" + target = "stop_billing" + return create_app(target, source, "cloudevent").test_client() + + +def test_receive_notification_to_stop_billing( + client: FlaskClient, + cloud_event_budget_alert: CloudEvent, + capsys: pytest.CaptureFixture[str] +) -> None: + headers, data = to_structured(cloud_event_budget_alert) + resp = client.post("/", headers=headers, data=data) + + captured = capsys.readouterr() + + assert resp.status_code == 200 + assert resp.data == b"OK" + + assert "Getting billing info for project" in captured.out + assert "Disabling billing for project" in captured.out + assert "Billing disabled. (Simulated)" in captured.out