Browse Source

Feat/compliance (#14982)

tags/1.0.1
Xiyuan Chen 7 months ago
parent
commit
9b9d14c2c4
No account linked to committer's email address

+ 1
- 1
api/controllers/console/__init__.py View File

from .auth import activate, data_source_bearer_auth, data_source_oauth, forgot_password, login, oauth from .auth import activate, data_source_bearer_auth, data_source_oauth, forgot_password, login, oauth


# Import billing controllers # Import billing controllers
from .billing import billing
from .billing import billing, compliance


# Import datasets controllers # Import datasets controllers
from .datasets import ( from .datasets import (

+ 35
- 0
api/controllers/console/billing/compliance.py View File

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")

+ 6
- 0
api/controllers/console/error.py View File

"This email account has been deleted within the past 30 days" "This email account has been deleted within the past 30 days"
"and is temporarily unavailable for new account registration." "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

+ 29
- 0
api/services/billing_service.py View File

from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fixed from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fixed


from extensions.ext_database import db from extensions.ext_database import db
from libs.helper import RateLimiter
from models.account import TenantAccountJoin, TenantAccountRole from models.account import TenantAccountJoin, TenantAccountRole




base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL") base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL")
secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY") 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 @classmethod
def get_info(cls, tenant_id: str): def get_info(cls, tenant_id: str):
params = {"tenant_id": tenant_id} params = {"tenant_id": tenant_id}
"""Update account deletion feedback.""" """Update account deletion feedback."""
json = {"email": email, "feedback": feedback} json = {"email": email, "feedback": feedback}
return cls._send_request("POST", "/account/delete-feedback", json=json) 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

Loading…
Cancel
Save