@cryptotaxi247 / CoPilot / commits / 1dde7bef

License cache (#498)

* Add LicenseCache model and migration for license caching functionality * Implement caching for license features and normalize API responses * precommit fixes

taylor_socfortress committed Sep 2, 2025 at 10:27 UTC 1dde7bef2ed7cb39d613a66abb23fcb7cf7630f5
3 files changed +477 -192
backend/alembic/versions/6796a7d001ad_add_license_cache_table.py new
+50
@@ -0,0 +1,50 @@
1 +"""Add license cache table
2 +
3 +Revision ID: 6796a7d001ad
4 +Revises: 8e3b57c8bbb2
5 +Create Date: 2025-09-02 09:58:22.242700
6 +
7 +"""
8 +from typing import Sequence
9 +from typing import Union
10 +
11 +import sqlalchemy as sa
12 +
13 +from alembic import op
14 +
15 +# revision identifiers, used by Alembic.
16 +revision: str = "6796a7d001ad"
17 +down_revision: Union[str, None] = "8e3b57c8bbb2"
18 +branch_labels: Union[str, Sequence[str], None] = None
19 +depends_on: Union[str, Sequence[str], None] = None
20 +
21 +
22 +def upgrade() -> None:
23 + # ### commands auto generated by Alembic - please adjust! ###
24 + op.create_table(
25 + "license_cache",
26 + sa.Column("id", sa.Integer(), nullable=False),
27 + sa.Column("license_key", sa.String(length=1024), nullable=False),
28 + sa.Column("feature_name", sa.String(length=256), nullable=False),
29 + sa.Column("is_enabled", sa.Boolean(), nullable=False),
30 + sa.Column("cached_at", sa.DateTime(), nullable=False),
31 + sa.Column("expires_at", sa.DateTime(), nullable=False),
32 + sa.Column("license_data", sa.String(length=5000), nullable=True),
33 + sa.PrimaryKeyConstraint("id"),
34 + )
35 + op.create_index(op.f("ix_license_cache_cached_at"), "license_cache", ["cached_at"], unique=False)
36 + op.create_index(op.f("ix_license_cache_expires_at"), "license_cache", ["expires_at"], unique=False)
37 + op.create_index(op.f("ix_license_cache_feature_name"), "license_cache", ["feature_name"], unique=False)
38 +
39 + # Create partial index on license_key (first 255 characters to avoid MySQL key length limit)
40 + op.execute("CREATE INDEX ix_license_cache_license_key ON license_cache (license_key(255))")
41 +
42 +
43 +def downgrade() -> None:
44 + # ### commands auto generated by Alembic - please adjust! ###
45 + op.drop_index("ix_license_cache_license_key", table_name="license_cache")
46 + op.drop_index(op.f("ix_license_cache_feature_name"), table_name="license_cache")
47 + op.drop_index(op.f("ix_license_cache_expires_at"), table_name="license_cache")
48 + op.drop_index(op.f("ix_license_cache_cached_at"), table_name="license_cache")
49 + op.drop_table("license_cache")
50 + # ### end Alembic commands ###
backend/app/db/universal_models.py
+11
@@ -239,6 +239,17 @@ class License(SQLModel, table=True):
239 company_name: str = Field(max_length=1024)
240
241
242 +class LicenseCache(SQLModel, table=True):
243 + __tablename__ = "license_cache"
244 + id: Optional[int] = Field(primary_key=True)
245 + license_key: str = Field(max_length=1024, index=True)
246 + feature_name: str = Field(max_length=256, index=True)
247 + is_enabled: bool = Field(default=False)
248 + cached_at: datetime = Field(default=datetime.utcnow, index=True)
249 + expires_at: datetime = Field(index=True)
250 + license_data: Optional[str] = Field(max_length=5000) # Store full license JSON as string for reference
251 +
252 +
253 class SchedulerJob(SQLModel, table=True):
254 id: str = Field(default=None, primary_key=True, nullable=False, max_length=255)
255 next_run_time: float = Field(sa_column=Column(Float(), index=True))
backend/app/middleware/license.py
+416 -192
@@ -1,5 +1,7 @@
1 +import json
2 import os
3 from datetime import datetime as dt
4 +from datetime import timedelta
5 from typing import Any
6 from typing import Dict
7 from typing import List
@@ -12,6 +14,7 @@ from fastapi import HTTPException
14 from loguru import logger
15 from pydantic import BaseModel
16 from pydantic import Field
17 +from sqlalchemy import delete
18 from sqlalchemy import select
19 from sqlalchemy.ext.asyncio import AsyncSession
20
@@ -19,6 +22,7 @@ from app.connectors.schema import UpdateConnector
22 from app.connectors.services import ConnectorServices
23 from app.db.db_session import get_db
24 from app.db.universal_models import License
25 +from app.db.universal_models import LicenseCache
26
27
28 class ThreatIntelRegisterRequest(BaseModel):
@@ -322,6 +326,28 @@ class RetrieveDockerCompose(BaseModel):
326
327 license_router = APIRouter()
328
329 +# Cache duration in hours
330 +CACHE_DURATION_HOURS = 1
331 +
332 +
333 +def normalize_api_response(response: Dict[str, Any]) -> Dict[str, Any]:
334 + """
335 + Normalize API responses to handle different response structures.
336 + Some endpoints return data wrapped in 'data' key, others don't.
337 +
338 + Args:
339 + response: Raw API response
340 +
341 + Returns:
342 + Normalized response with consistent structure
343 + """
344 + if "data" in response:
345 + # Response has data wrapper - return as is
346 + return response
347 + else:
348 + # Response doesn't have data wrapper - wrap it
349 + return {"data": response, "success": response.get("success", True), "message": response.get("message", "Success")}
350 +
351
352 async def check_if_license_exists(session: AsyncSession):
353 # Get the first row and raise HTTPException stating license already exists
@@ -377,6 +403,122 @@ async def get_license(session: AsyncSession) -> License:
403 return license
404
405
406 +async def get_cached_feature(session: AsyncSession, license_key: str, feature_name: str) -> Optional[LicenseCache]:
407 + """
408 + Get cached feature information if it exists and is not expired.
409 +
410 + Args:
411 + session: Database session
412 + license_key: The license key
413 + feature_name: The feature name to check
414 +
415 + Returns:
416 + LicenseCache object if valid cache exists, None otherwise
417 + """
418 + current_time = dt.utcnow()
419 +
420 + result = await session.execute(
421 + select(LicenseCache).where(
422 + LicenseCache.license_key == license_key,
423 + LicenseCache.feature_name == feature_name,
424 + LicenseCache.expires_at > current_time,
425 + ),
426 + )
427 +
428 + cached_feature = result.scalars().first()
429 +
430 + if cached_feature:
431 + logger.info(f"Found valid cache for feature '{feature_name}' (expires at {cached_feature.expires_at})")
432 + return cached_feature
433 + else:
434 + logger.info(f"No valid cache found for feature '{feature_name}'")
435 + return None
436 +
437 +
438 +async def cache_license_features(session: AsyncSession, license_key: str, license_data: Dict[str, Any]) -> None:
439 + """
440 + Cache license features from license verification response.
441 +
442 + Args:
443 + session: Database session
444 + license_key: The license key
445 + license_data: The license verification response data
446 + """
447 + try:
448 + # Clear existing cache for this license key
449 + await session.execute(delete(LicenseCache).where(LicenseCache.license_key == license_key))
450 +
451 + current_time = dt.utcnow()
452 + expires_at = current_time + timedelta(hours=CACHE_DURATION_HOURS)
453 +
454 + # Store the full license data as JSON string for reference
455 + license_json = json.dumps(license_data)
456 +
457 + # Extract features from dataObjects - handle both response formats
458 + if "data" in license_data and "license" in license_data["data"]:
459 + # Wrapped format: {"data": {"license": {"dataObjects": [...]}}}
460 + data_objects = license_data["data"]["license"].get("dataObjects", [])
461 + elif "license" in license_data:
462 + # Direct format: {"license": {"dataObjects": [...]}}
463 + data_objects = license_data["license"].get("dataObjects", [])
464 + elif "dataObjects" in license_data:
465 + # Bare format: {"dataObjects": [...]}
466 + data_objects = license_data.get("dataObjects", [])
467 + else:
468 + logger.warning("Could not find dataObjects in license response")
469 + data_objects = []
470 +
471 + # Track which features we've processed
472 + processed_features = set()
473 +
474 + for data_object in data_objects:
475 + feature_name = data_object.get("name")
476 + is_enabled = data_object.get("intValue") == 1
477 +
478 + if feature_name:
479 + cache_entry = LicenseCache(
480 + license_key=license_key,
481 + feature_name=feature_name,
482 + is_enabled=is_enabled,
483 + cached_at=current_time,
484 + expires_at=expires_at,
485 + license_data=license_json,
486 + )
487 +
488 + session.add(cache_entry)
489 + processed_features.add(feature_name)
490 +
491 + logger.info(f"Cached feature '{feature_name}': {'enabled' if is_enabled else 'disabled'}")
492 +
493 + await session.commit()
494 + logger.info(f"Successfully cached {len(processed_features)} features for license {license_key[:8]}...")
495 +
496 + except Exception as e:
497 + logger.error(f"Error caching license features: {str(e)}")
498 + await session.rollback()
499 + raise
500 +
501 +
502 +async def invalidate_license_cache(session: AsyncSession, license_key: str) -> None:
503 + """
504 + Invalidate (delete) all cached entries for a specific license key.
505 +
506 + Args:
507 + session: Database session
508 + license_key: The license key to invalidate cache for
509 + """
510 + try:
511 + result = await session.execute(delete(LicenseCache).where(LicenseCache.license_key == license_key))
512 + deleted_count = result.rowcount
513 + await session.commit()
514 +
515 + logger.info(f"Invalidated {deleted_count} cache entries for license {license_key[:8]}...")
516 +
517 + except Exception as e:
518 + logger.error(f"Error invalidating license cache: {str(e)}")
519 + await session.rollback()
520 +
521 +
522 def is_license_expired(license: dict) -> bool:
523 """
524 Check if a license is expired.
@@ -395,28 +537,72 @@ def is_license_expired(license: dict) -> bool:
537 async def is_feature_enabled(feature_name: str, session: AsyncSession, message: str = None) -> bool:
538 """
539 Check if a feature is enabled in a license.
540 + Uses cache first, falls back to API if cache miss or expired.
541
542 Args:
400 - license (License): The license to check.
543 feature_name (str): The feature name to check.
544 session (AsyncSession): The database session.
545 + message (str, optional): Custom error message.
546
547 Returns:
548 bool: True if the feature is enabled, False otherwise.
549 """
550 license = await get_license(session)
408 - result = await send_post_request("verify-license", data={"license_key": license.license_key})
409 - for data_object in result["data"]["license"]["dataObjects"]:
410 - if data_object["name"] == feature_name and data_object["intValue"] == 1:
551 +
552 + # Check cache first
553 + cached_feature = await get_cached_feature(session, license.license_key, feature_name)
554 +
555 + if cached_feature:
556 + logger.info(f"Using cached result for feature '{feature_name}': {'enabled' if cached_feature.is_enabled else 'disabled'}")
557 +
558 + if cached_feature.is_enabled:
559 return True
560 + else:
561 + # Feature is disabled according to cache
562 + if message:
563 + raise HTTPException(status_code=400, detail=message)
564 + raise HTTPException(
565 + status_code=400,
566 + detail=f"Feature is not enabled. You must purchase the {feature_name} license to use this feature.",
567 + )
568
413 - if message:
414 - raise HTTPException(status_code=400, detail=message)
569 + # Cache miss - fetch from API
570 + logger.info(f"Cache miss for feature '{feature_name}', fetching from API")
571
416 - raise HTTPException(
417 - status_code=400,
418 - detail=f"Feature is not enabled. You must purchase the {feature_name} license to use this feature.",
419 - )
572 + try:
573 + result = await send_post_request("verify-license", data={"license_key": license.license_key})
574 +
575 + # Normalize response format
576 + normalized_result = normalize_api_response(result)
577 +
578 + # Cache the results
579 + await cache_license_features(session, license.license_key, normalized_result)
580 +
581 + # Check if feature is enabled
582 + data_objects = normalized_result["data"]["license"].get("dataObjects", [])
583 + for data_object in data_objects:
584 + if data_object["name"] == feature_name and data_object["intValue"] == 1:
585 + logger.info(f"Feature '{feature_name}' is enabled (from API)")
586 + return True
587 +
588 + # Feature not found or not enabled
589 + logger.info(f"Feature '{feature_name}' is not enabled (from API)")
590 +
591 + if message:
592 + raise HTTPException(status_code=400, detail=message)
593 +
594 + raise HTTPException(
595 + status_code=400,
596 + detail=f"Feature is not enabled. You must purchase the {feature_name} license to use this feature.",
597 + )
598 +
599 + except HTTPException:
600 + # Re-raise HTTP exceptions (like feature not enabled)
601 + raise
602 + except Exception as e:
603 + logger.error(f"Error verifying license from API: {str(e)}")
604 + # If API fails, we can't verify - raise an error
605 + raise HTTPException(status_code=500, detail=f"Unable to verify license: {str(e)}")
606
607
608 async def send_get_request(endpoint: str) -> Dict[str, Any]:
@@ -437,6 +623,7 @@ async def send_get_request(endpoint: str) -> Dict[str, Any]:
623 "Content-Type": "application/json",
624 "module-version": "1.0",
625 }
626 + logger.info(f"Request headers: {HEADERS}")
627 response = requests.get(
628 f"https://license.socfortress.co/{endpoint}",
629 headers=HEADERS,
@@ -444,17 +631,9 @@ async def send_get_request(endpoint: str) -> Dict[str, Any]:
631 )
632
633 if response.status_code == 204:
447 - return {
448 - "data": None,
449 - "success": True,
450 - "message": "Successfully completed request with no content",
451 - }
634 + return {"success": True, "message": "No content"}
635 else:
453 - return {
454 - "data": response.json(),
455 - "success": False if response.status_code >= 400 else True,
456 - "message": "Successfully retrieved data",
457 - }
636 + return response.json()
637 except Exception as e:
638 logger.error(f"Failed to send GET request to {endpoint} with error: {e}")
639 raise HTTPException(
@@ -476,15 +655,31 @@ async def get_subscription_catalog():
655 dict: A dictionary containing the subscription catalog.
656 """
657 try:
479 - results = await send_get_request("features")
658 + result = await send_get_request("features")
659 + normalized_result = normalize_api_response(result)
660 +
661 + # Handle different response structures for subscription features
662 + features = []
663 + if "data" in normalized_result and "features" in normalized_result["data"]:
664 + features = normalized_result["data"]["features"]
665 + elif "features" in normalized_result:
666 + features = normalized_result["features"]
667 + else:
668 + logger.warning(f"No features found in response: {normalized_result}")
669 + features = []
670 +
671 return GetSubscriptionCatalogFeaturesResponse(
481 - features=results["data"]["features"],
482 - success=results["success"],
483 - message=results["message"],
672 + features=features,
673 + success=normalized_result.get("success", True),
674 + message=normalized_result.get("message", "Subscription features retrieved successfully"),
675 )
676 except Exception as e:
486 - logger.error(e)
487 - raise HTTPException(status_code=400, detail="Failed to get subscription features")
677 + logger.error(f"Error getting subscription catalog: {str(e)}")
678 + return GetSubscriptionCatalogFeaturesResponse(
679 + features=[],
680 + success=False,
681 + message=f"Error getting subscription catalog: {str(e)}",
682 + )
683
684
685 @license_router.post(
@@ -510,28 +705,29 @@ async def retrieve_license_by_email(request: GetLicenseByEmailRequest, session:
705 return GetLicenseResponse(
706 license_key=existing_license.license_key,
707 success=True,
513 - message="License retrieved successfully",
708 + message="License already exists in database",
709 )
710
711 results = await send_post_request("retrieve-license-by-email", data={"email": request.email})
517 - logger.info(f"Results: {results}")
518 - if results["data"]["success"] is False:
519 - raise HTTPException(status_code=400, detail=f"Failed to retrieve license by email: {results['data']['message']}")
712 + normalized_results = normalize_api_response(results)
713 + logger.info(f"Results: {normalized_results}")
714 + if normalized_results["data"]["success"] is False:
715 + raise HTTPException(status_code=400, detail=normalized_results["data"]["message"])
716
717 # Add the license to the database
718 await add_license_to_db(
719 session,
524 - results["data"]["license"]["key"],
720 + normalized_results["data"]["license"]["key"],
721 AddLicenseToDB(
526 - customer_email=results["data"]["license"]["customer"]["email"],
527 - customer_name=results["data"]["license"]["customer"]["name"],
528 - company_name=results["data"]["license"]["customer"]["companyName"],
722 + customer_email=normalized_results["data"]["license"]["customer"]["email"],
723 + customer_name=normalized_results["data"]["license"]["customer"]["name"],
724 + company_name=normalized_results["data"]["license"]["customer"]["companyName"],
725 ),
726 )
727 return GetLicenseResponse(
532 - license_key=results["data"]["license"]["key"],
533 - success=results["data"]["success"],
534 - message=results["data"]["message"],
728 + license_key=normalized_results["data"]["license"]["key"],
729 + success=normalized_results["data"]["success"],
730 + message=normalized_results["data"]["message"],
731 )
732
733
@@ -560,13 +756,14 @@ async def create_checkout_session(request: FeatureSubscriptionRequest):
756 "company_name": request.company_name,
757 },
758 )
563 - logger.info(f"Results: {results}")
564 - if results["data"]["success"] is False:
565 - raise HTTPException(status_code=400, detail=f"Failed to create checkout session: {results['data']['message']}")
759 + normalized_results = normalize_api_response(results)
760 + logger.info(f"Results: {normalized_results}")
761 + if normalized_results["data"]["success"] is False:
762 + raise HTTPException(status_code=400, detail=normalized_results["data"]["message"])
763 return CheckoutSessionResponse(
567 - session=results["data"]["session"],
568 - success=results["data"]["success"],
569 - message=results["data"]["message"],
764 + session=normalized_results["data"]["session"],
765 + success=normalized_results["data"]["success"],
766 + message=normalized_results["data"]["message"],
767 )
768
769
@@ -587,32 +784,37 @@ async def create_trial_license_key(request: TrialLicenseRequest, session: AsyncS
784 LicenseVerificationResponse: A Pydantic model containing the verification status and message.
785 """
786 await check_if_license_exists(session)
787 +
788 results = await send_post_request(
789 "trial-license",
790 data={
791 + "period": request.period,
792 "email": request.email,
793 "feature_name": request.feature_name,
794 "customer_name": request.customer_name,
596 - "period": request.period,
795 "company_name": request.company_name,
796 },
797 )
600 - logger.info(f"Results: {results}")
601 - if results["data"]["success"] is False:
602 - raise HTTPException(status_code=400, detail=f"Failed to create trial license: {results['data']['message']}")
798 + normalized_results = normalize_api_response(results)
799 + logger.info(f"Results: {normalized_results}")
800 + if normalized_results["data"]["success"] is False:
801 + raise HTTPException(status_code=400, detail=normalized_results["data"]["message"])
802 +
803 + # Add the license to the database
804 await add_license_to_db(
805 session,
605 - results["data"]["license_key"],
806 + normalized_results["data"]["license_key"],
807 AddLicenseToDB(
808 customer_email=request.email,
809 customer_name=request.customer_name,
810 company_name=request.company_name,
811 ),
812 )
813 +
814 return TrialLicenseResponse(
613 - license_key=results["data"]["license_key"],
614 - success=results["data"]["success"],
615 - message=results["data"]["message"],
815 + license_key=normalized_results["data"]["license_key"],
816 + success=normalized_results["data"]["success"],
817 + message=normalized_results["data"]["message"],
818 )
819
820
@@ -622,15 +824,6 @@ async def create_trial_license_key(request: TrialLicenseRequest, session: AsyncS
824 response_model=CancelSubscriptionResponse,
825 )
826 async def cancel_subscription(request: CancelSubscriptionRequest) -> CancelSubscriptionResponse:
625 - """
626 - Cancel a subscription.
627 -
628 - Args:
629 - request (CancelSubscriptionRequest): The request containing the customer email, subscription price id, and feature name.
630 -
631 - Returns:
632 - dict: A dictionary containing the cancellation status.
633 - """
827 results = await send_post_request(
828 "cancel-subscription",
829 data={
@@ -639,12 +832,13 @@ async def cancel_subscription(request: CancelSubscriptionRequest) -> CancelSubsc
832 "feature_name": request.feature_name,
833 },
834 )
642 - logger.info(f"Results: {results}")
643 - if results["data"]["success"] is False:
644 - raise HTTPException(status_code=400, detail=f"Failed to cancel subscription: {results['data']['message']}")
835 + normalized_results = normalize_api_response(results)
836 + logger.info(f"Results: {normalized_results}")
837 + if normalized_results["data"]["success"] is False:
838 + raise HTTPException(status_code=400, detail=normalized_results["data"]["message"])
839 return CancelSubscriptionResponse(
646 - success=results["data"]["success"],
647 - message=results["data"]["message"],
840 + success=normalized_results["data"]["success"],
841 + message=normalized_results["data"]["message"],
842 )
843
844
@@ -654,24 +848,67 @@ async def cancel_subscription(request: CancelSubscriptionRequest) -> CancelSubsc
848 description="Verify a license key",
849 )
850 async def verify_license_key(session: AsyncSession = Depends(get_db)) -> VerifyLicenseResponse:
657 - """ "
658 - Verify a license key.
851 + license = await get_license(session)
852
660 - Args:
661 - license_key (str): The license key to verify.
853 + # Check if we have a recent cache entry for any feature of this license
854 + current_time = dt.utcnow()
855 + result = await session.execute(
856 + select(LicenseCache).where(LicenseCache.license_key == license.license_key, LicenseCache.expires_at > current_time).limit(1),
857 + )
858
663 - Returns:
664 - LicenseVerificationResponse: A Pydantic model containing the verification status and message.
665 - """
666 - license = await get_license(session)
667 - try:
668 - result = await send_post_request("verify-license", data={"license_key": license.license_key})
669 - if is_license_expired(result):
670 - raise HTTPException(status_code=400, detail="License is expired")
671 - return VerifyLicenseResponse(license=result["data"]["license"], success=True, message="License verified successfully")
672 - except Exception as e:
673 - logger.error(e)
674 - raise HTTPException(status_code=400, detail="License verification failed")
859 + cached_entry = result.scalars().first()
860 +
861 + if cached_entry and cached_entry.license_data:
862 + logger.info("Using cached license verification data")
863 + try:
864 + license_data = json.loads(cached_entry.license_data)
865 + # Handle both wrapped and unwrapped formats
866 + if "data" in license_data and "license" in license_data["data"]:
867 + license_obj = license_data["data"]["license"]
868 + success = license_data["data"]["success"]
869 + elif "license" in license_data:
870 + license_obj = license_data["license"]
871 + success = license_data.get("success", True)
872 + else:
873 + raise ValueError("Invalid cached license data format")
874 +
875 + return VerifyLicenseResponse(
876 + license=license_obj,
877 + success=success,
878 + message="License verified successfully (from cache)",
879 + )
880 + except (json.JSONDecodeError, KeyError, ValueError) as e:
881 + logger.warning(f"Error parsing cached license data: {e}, falling back to API")
882 +
883 + # No valid cache, fetch from API
884 + logger.info("No valid cache found, verifying license via API")
885 + results = await send_post_request("verify-license", data={"license_key": license.license_key})
886 +
887 + # Normalize response format
888 + normalized_results = normalize_api_response(results)
889 +
890 + # Cache the results
891 + await cache_license_features(session, license.license_key, normalized_results)
892 +
893 + logger.info(f"Results: {normalized_results}")
894 + if not normalized_results.get("success", True):
895 + raise HTTPException(status_code=400, detail=normalized_results.get("message", "License verification failed"))
896 +
897 + # Handle both response formats for return value
898 + if "data" in normalized_results and "license" in normalized_results["data"]:
899 + license_obj = normalized_results["data"]["license"]
900 + success = normalized_results["data"]["success"]
901 + message = normalized_results["data"]["message"]
902 + else:
903 + license_obj = normalized_results["license"]
904 + success = normalized_results.get("success", True)
905 + message = normalized_results.get("message", "License verified successfully")
906 +
907 + return VerifyLicenseResponse(
908 + license=license_obj,
909 + success=success,
910 + message=message,
911 + )
912
913
914 @license_router.get(
@@ -679,17 +916,12 @@ async def verify_license_key(session: AsyncSession = Depends(get_db)) -> VerifyL
916 description="Get a license",
917 )
918 async def get_license_key(session: AsyncSession = Depends(get_db)) -> GetLicenseResponse:
682 - """ "
683 - Get a license key.
684 -
685 - Args:
686 - license_key (str): The license key to verify.
687 -
688 - Returns:
689 - LicenseVerificationResponse: A Pydantic model containing the verification status and message.
690 - """
919 license = await get_license(session)
692 - return GetLicenseResponse(license_key=license.license_key, success=True, message="License retrieved successfully")
920 + return GetLicenseResponse(
921 + license_key=license.license_key,
922 + success=True,
923 + message="License retrieved successfully",
924 + )
925
926
927 async def send_post_request(endpoint: str, data: Dict[str, Any] = None) -> Dict[str, Any]:
@@ -699,7 +931,6 @@ async def send_post_request(endpoint: str, data: Dict[str, Any] = None) -> Dict[
931 Args:
932 endpoint (str): The endpoint to send the POST request to.
933 data (Dict[str, Any]): The data to send with the POST request.
702 - connector_name (str, optional): The name of the connector to use. Defaults to "Shuffle".
934
935 Returns:
936 Dict[str, Any]: The response from the POST request.
@@ -719,17 +950,10 @@ async def send_post_request(endpoint: str, data: Dict[str, Any] = None) -> Dict[
950 verify=False,
951 )
952
722 - if response.status_code == 200:
723 - return {
724 - "data": response.json(),
725 - "success": True,
726 - "message": "Successfully retrieved data",
727 - }
953 + if response.status_code == 204:
954 + return {"success": True, "message": "No content"}
955 else:
729 - return {
730 - "success": False,
731 - "message": f"Failed to send POST request to {endpoint}",
732 - }
956 + return response.json()
957 except Exception as e:
958 logger.error(f"Failed to send POST request to {endpoint} with error: {e}")
959 raise HTTPException(
@@ -744,28 +968,22 @@ async def send_post_request(endpoint: str, data: Dict[str, Any] = None) -> Dict[
968 description="Check if a feature is enabled in a license",
969 )
970 async def is_feature_enabled_route(feature_name: str, session: AsyncSession = Depends(get_db)) -> IsFeatureEnabledResponse:
747 - """
748 - Check if a feature is enabled in a license.
749 -
750 - Args:
751 - feature_name (str): The feature name to check.
752 - session (AsyncSession, optional): The database session. Defaults to Depends(get_db).
753 -
754 - Returns:
755 - bool: True if the feature is enabled, False otherwise.
756 - """
757 - if await is_feature_enabled(feature_name, session):
971 + try:
972 + await is_feature_enabled(feature_name, session)
973 return IsFeatureEnabledResponse(
974 enabled=True,
975 success=True,
761 - message="Feature is enabled",
762 - )
763 - else:
764 - return IsFeatureEnabledResponse(
765 - enabled=False,
766 - success=True,
767 - message="Feature is not enabled",
976 + message=f"Feature '{feature_name}' is enabled",
977 )
978 + except HTTPException as e:
979 + if e.status_code == 400: # Feature not enabled
980 + return IsFeatureEnabledResponse(
981 + enabled=False,
982 + success=True,
983 + message=e.detail,
984 + )
985 + else:
986 + raise
987
988
989 @license_router.get(
@@ -774,26 +992,54 @@ async def is_feature_enabled_route(feature_name: str, session: AsyncSession = De
992 description="Get license features",
993 )
994 async def get_license_features(session: AsyncSession = Depends(get_db)) -> GetLicenseFeaturesResponse:
777 - """
778 - Get the features enabled in a license.
995 + license = await get_license(session)
996
780 - Args:
781 - session (AsyncSession, optional): The database session. Defaults to Depends(get_db).
997 + # Try to get features from cache first
998 + current_time = dt.utcnow()
999 + result = await session.execute(
1000 + select(LicenseCache).where(
1001 + LicenseCache.license_key == license.license_key,
1002 + LicenseCache.expires_at > current_time,
1003 + LicenseCache.is_enabled == True,
1004 + ),
1005 + )
1006
783 - Returns:
784 - dict: A dictionary containing the features enabled in the license.
785 - """
786 - license = await get_license(session)
787 - try:
788 - results = await send_post_request("license-features", data={"license_key": license.license_key})
1007 + cached_features = result.scalars().all()
1008 +
1009 + if cached_features:
1010 + logger.info(f"Using cached license features ({len(cached_features)} features)")
1011 + features = [cache.feature_name for cache in cached_features]
1012 return GetLicenseFeaturesResponse(
790 - features=results["data"]["features"],
791 - success=results["success"],
792 - message=results["message"],
1013 + features=features,
1014 + success=True,
1015 + message="License features retrieved successfully (from cache)",
1016 )
794 - except Exception as e:
795 - logger.error(e)
796 - raise HTTPException(status_code=400, detail="Failed to get license features")
1017 +
1018 + # No cache, fetch from API
1019 + logger.info("No cached features found, fetching from API")
1020 + results = await send_post_request("license-features", data={"license_key": license.license_key})
1021 +
1022 + # This endpoint returns features directly without data wrapper
1023 + logger.info(f"Results: {results}")
1024 + if not results.get("success", True):
1025 + raise HTTPException(status_code=400, detail=results.get("message", "Failed to get license features"))
1026 +
1027 + # For license-features endpoint, create a fake normalized response to cache the features
1028 + fake_license_data = {
1029 + "data": {
1030 + "license": {"dataObjects": [{"name": feature, "intValue": 1} for feature in results.get("features", [])]},
1031 + "success": True,
1032 + },
1033 + }
1034 +
1035 + # Cache the results
1036 + await cache_license_features(session, license.license_key, fake_license_data)
1037 +
1038 + return GetLicenseFeaturesResponse(
1039 + features=results.get("features", []),
1040 + success=results.get("success", True),
1041 + message=results.get("message", "License features retrieved successfully"),
1042 + )
1043
1044
1045 @license_router.post(
@@ -801,39 +1047,17 @@ async def get_license_features(session: AsyncSession = Depends(get_db)) -> GetLi
1047 description="Replace a license",
1048 )
1049 async def replace_license_in_db(request: ReplaceLicenseRequest, session: AsyncSession = Depends(get_db)):
804 - """
805 - Replace a license in the database.
1050 + license = await get_license(session)
1051
807 - Args:
808 - request (ReplaceLicenseRequest): The request containing the license key to replace.
809 - session (AsyncSession, optional): The database session. Defaults to Depends(get_db).
1052 + # Invalidate cache for old license
1053 + await invalidate_license_cache(session, license.license_key)
1054
811 - Returns:
812 - LicenseVerificationResponse: A Pydantic model containing the verification status and message.
813 - """
814 - try:
815 - # Update the license in the database
816 - result = await session.execute(select(License))
817 - license = result.scalars().first()
818 - if not license:
819 - # Verify the license key
820 - license_data = await send_post_request("verify-license", data={"license_key": request.license_key})
821 - if is_license_expired(license_data):
822 - raise HTTPException(status_code=400, detail="License is expired")
823 - # Create a new License object with the data from the dictionary
824 - license = License(
825 - license_key=license_data["data"]["license"]["key"],
826 - customer_name=license_data["data"]["license"]["customer"]["name"],
827 - customer_email=license_data["data"]["license"]["customer"]["email"],
828 - company_name=license_data["data"]["license"]["customer"]["companyName"],
829 - )
830 - session.add(license)
831 - license.license_key = request.license_key
832 - await session.commit()
833 - return {"message": "License replaced successfully", "success": True}
834 - except Exception as e:
835 - logger.error(e)
836 - raise HTTPException(status_code=400, detail="License replacement failed")
1055 + # Update license key
1056 + license.license_key = request.license_key
1057 + await session.commit()
1058 +
1059 + logger.info("License replaced successfully")
1060 + return {"success": True, "message": "License replaced successfully"}
1061
1062
1063 @license_router.post(
@@ -842,34 +1066,34 @@ async def replace_license_in_db(request: ReplaceLicenseRequest, session: AsyncSe
1066 description="Retrieve Docker Compose for features enabled",
1067 )
1068 async def retrieve_docker_compose(session: AsyncSession = Depends(get_db)) -> RetrieveDockerCompose:
845 - """
846 - Retrieve the Docker Compose for features enabled in the license.
1069 + license = await get_license(session)
1070 + results = await send_post_request("retrieve-docker-compose", data={"license_key": license.license_key})
1071
848 - Args:
849 - session (AsyncSession, optional): The database session. Defaults to Depends(get_db).
1072 + # This endpoint returns data directly without wrapper
1073 + logger.info(f"Results: {results}")
1074 + if not results.get("success", True):
1075 + raise HTTPException(status_code=400, detail=results.get("message", "Failed to retrieve Docker Compose"))
1076
851 - Returns:
852 - RetrieveDockerCompose: A Pydantic model containing the Docker Compose for features enabled.
1077 + return RetrieveDockerCompose(
1078 + docker_compose=results.get("docker_compose", ""),
1079 + success=results.get("success", True),
1080 + message=results.get("message", "Docker Compose retrieved successfully"),
1081 + )
1082 +
1083 +
1084 +@license_router.post(
1085 + "/invalidate_cache",
1086 + description="Manually invalidate license cache",
1087 +)
1088 +async def invalidate_cache_route(session: AsyncSession = Depends(get_db)):
1089 """
854 - try:
855 - license = await get_license(session)
856 - if license.license_key:
857 - results = await send_post_request("retrieve-docker-compose", data={"license_key": license.license_key})
858 - logger.info(f"Results: {results}")
859 - return RetrieveDockerCompose(
860 - docker_compose=results["data"]["docker_compose"],
861 - success=results["success"],
862 - message=results["message"],
863 - )
864 - else:
865 - return RetrieveDockerCompose(
866 - docker_compose="",
867 - success=False,
868 - message="License key not found",
869 - )
870 - except Exception as e:
871 - logger.error(e)
872 - raise HTTPException(status_code=400, detail="Failed to retrieve Docker Compose")
1090 + Manually invalidate the license cache.
1091 + Useful for testing or when you want to force a fresh license check.
1092 + """
1093 + license = await get_license(session)
1094 + await invalidate_license_cache(session, license.license_key)
1095 +
1096 + return {"success": True, "message": "License cache invalidated successfully"}
1097
1098
1099 def create_headers(request: ThreatIntelRegisterRequest) -> Dict[str, str]: