From ed7d7a74ea475400a9f4b13e148cbd16cac8cc17 Mon Sep 17 00:00:00 2001
From: Xiyuan Chen <52963600+GareArc@users.noreply.github.com>
Date: Sun, 16 Feb 2025 20:31:52 -0500
Subject: [PATCH] Feat/compliance (#13548)

---
 api/controllers/console/__init__.py           |  2 +-
 api/controllers/console/billing/compliance.py | 35 +++++++++++++++++++
 api/controllers/console/error.py              |  6 ++++
 api/services/billing_service.py               | 29 +++++++++++++++
 4 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 api/controllers/console/billing/compliance.py

diff --git a/api/controllers/console/__init__.py b/api/controllers/console/__init__.py
index cb6b0d097b..a630f16be7 100644
--- a/api/controllers/console/__init__.py
+++ b/api/controllers/console/__init__.py
@@ -70,7 +70,7 @@ from .app import (
 from .auth import activate, data_source_bearer_auth, data_source_oauth, forgot_password, login, oauth
 
 # Import billing controllers
-from .billing import billing
+from .billing import billing, compliance
 
 # Import datasets controllers
 from .datasets import (
diff --git a/api/controllers/console/billing/compliance.py b/api/controllers/console/billing/compliance.py
new file mode 100644
index 0000000000..6d5d668709
--- /dev/null
+++ b/api/controllers/console/billing/compliance.py
@@ -0,0 +1,35 @@
+from flask import request
+from flask_login import current_user  # type: ignore
+from flask_restful import Resource, reqparse  # type: ignore
+
+from libs.helper import extract_remote_ip
+from libs.login import login_required
+from services.billing_service import BillingService
+
+from .. import api
+from ..wraps import account_initialization_required, only_edition_cloud, setup_required
+
+
+class ComplianceApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @only_edition_cloud
+    def get(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument("doc_name", type=str, required=True, location="args")
+        args = parser.parse_args()
+
+        ip_address = extract_remote_ip(request)
+        device_info = request.headers.get("User-Agent", "Unknown device")
+
+        return BillingService.get_compliance_download_link(
+            doc_name=args.doc_name,
+            account_id=current_user.id,
+            tenant_id=current_user.current_tenant_id,
+            ip=ip_address,
+            device_info=device_info,
+        )
+
+
+api.add_resource(ComplianceApi, "/compliance/download")
diff --git a/api/controllers/console/error.py b/api/controllers/console/error.py
index ee87138a44..bd4ae9dc7f 100644
--- a/api/controllers/console/error.py
+++ b/api/controllers/console/error.py
@@ -101,3 +101,9 @@ class AccountInFreezeError(BaseHTTPException):
         "This email account has been deleted within the past 30 days"
         "and is temporarily unavailable for new account registration."
     )
+
+
+class CompilanceRateLimitError(BaseHTTPException):
+    error_code = "compilance_rate_limit"
+    description = "Rate limit exceeded for downloading compliance report."
+    code = 429
diff --git a/api/services/billing_service.py b/api/services/billing_service.py
index 0d50a2aa8c..ad141035cc 100644
--- a/api/services/billing_service.py
+++ b/api/services/billing_service.py
@@ -5,6 +5,7 @@ import httpx
 from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fixed
 
 from extensions.ext_database import db
+from libs.helper import RateLimiter
 from models.account import TenantAccountJoin, TenantAccountRole
 
 
@@ -12,6 +13,8 @@ class BillingService:
     base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL")
     secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY")
 
+    compliance_download_rate_limiter = RateLimiter("compliance_download_rate_limiter", 4, 60)
+
     @classmethod
     def get_info(cls, tenant_id: str):
         params = {"tenant_id": tenant_id}
@@ -91,3 +94,29 @@ class BillingService:
         """Update account deletion feedback."""
         json = {"email": email, "feedback": feedback}
         return cls._send_request("POST", "/account/delete-feedback", json=json)
+
+    @classmethod
+    def get_compliance_download_link(
+        cls,
+        doc_name: str,
+        account_id: str,
+        tenant_id: str,
+        ip: str,
+        device_info: str,
+    ):
+        limiter_key = f"{account_id}:{tenant_id}"
+        if cls.compliance_download_rate_limiter.is_rate_limited(limiter_key):
+            from controllers.console.error import CompilanceRateLimitError
+
+            raise CompilanceRateLimitError()
+
+        json = {
+            "doc_name": doc_name,
+            "account_id": account_id,
+            "tenant_id": tenant_id,
+            "ip_address": ip,
+            "device_info": device_info,
+        }
+        res = cls._send_request("POST", "/compliance/download", json=json)
+        cls.compliance_download_rate_limiter.increment_rate_limit(limiter_key)
+        return res