dify

dify/api/controllers/console/billing/billing.py

@@ -0,0 +1,80 @@
+import base64
+
+from flask_restx import Resource, fields, reqparse
+from werkzeug.exceptions import BadRequest
+
+from controllers.console import console_ns
+from controllers.console.wraps import account_initialization_required, only_edition_cloud, setup_required
+from enums.cloud_plan import CloudPlan
+from libs.login import current_account_with_tenant, login_required
+from services.billing_service import BillingService
+
+
+@console_ns.route("/billing/subscription")
+class Subscription(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @only_edition_cloud
+    def get(self):
+        current_user, current_tenant_id = current_account_with_tenant()
+        parser = (
+            reqparse.RequestParser()
+            .add_argument(
+                "plan",
+                type=str,
+                required=True,
+                location="args",
+                choices=[CloudPlan.PROFESSIONAL, CloudPlan.TEAM],
+            )
+            .add_argument("interval", type=str, required=True, location="args", choices=["month", "year"])
+        )
+        args = parser.parse_args()
+        BillingService.is_tenant_owner_or_admin(current_user)
+        return BillingService.get_subscription(args["plan"], args["interval"], current_user.email, current_tenant_id)
+
+
+@console_ns.route("/billing/invoices")
+class Invoices(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @only_edition_cloud
+    def get(self):
+        current_user, current_tenant_id = current_account_with_tenant()
+        BillingService.is_tenant_owner_or_admin(current_user)
+        return BillingService.get_invoices(current_user.email, current_tenant_id)
+
+
+@console_ns.route("/billing/partners/<string:partner_key>/tenants")
+class PartnerTenants(Resource):
+    @console_ns.doc("sync_partner_tenants_bindings")
+    @console_ns.doc(description="Sync partner tenants bindings")
+    @console_ns.doc(params={"partner_key": "Partner key"})
+    @console_ns.expect(
+        console_ns.model(
+            "SyncPartnerTenantsBindingsRequest",
+            {"click_id": fields.String(required=True, description="Click Id from partner referral link")},
+        )
+    )
+    @console_ns.response(200, "Tenants synced to partner successfully")
+    @console_ns.response(400, "Invalid partner information")
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @only_edition_cloud
+    def put(self, partner_key: str):
+        current_user, _ = current_account_with_tenant()
+        parser = reqparse.RequestParser().add_argument("click_id", required=True, type=str, location="json")
+        args = parser.parse_args()
+
+        try:
+            click_id = args["click_id"]
+            decoded_partner_key = base64.b64decode(partner_key).decode("utf-8")
+        except Exception:
+            raise BadRequest("Invalid partner_key")
+
+        if not click_id or not decoded_partner_key or not current_user.id:
+            raise BadRequest("Invalid partner information")
+
+        return BillingService.sync_partner_tenants_bindings(current_user.id, decoded_partner_key, click_id)

dify/api/controllers/console/billing/compliance.py

@@ -0,0 +1,31 @@
+from flask import request
+from flask_restx import Resource, reqparse
+
+from libs.helper import extract_remote_ip
+from libs.login import current_account_with_tenant, login_required
+from services.billing_service import BillingService
+
+from .. import console_ns
+from ..wraps import account_initialization_required, only_edition_cloud, setup_required
+
+
+@console_ns.route("/compliance/download")
+class ComplianceApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @only_edition_cloud
+    def get(self):
+        current_user, current_tenant_id = current_account_with_tenant()
+        parser = reqparse.RequestParser().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_tenant_id,
+            ip=ip_address,
+            device_info=device_info,
+        )