1
import os
2
from datetime import datetime as dt
3
-from enum import Enum
3
from typing import Any
4
from typing import Dict
5
from typing import List
9
from fastapi import APIRouter
10
from fastapi import Depends
11
from fastapi import HTTPException
13
-from licensing.methods import Data
14
-from licensing.methods import Helpers
15
-from licensing.methods import Key
16
-
17
-# from licensing.models import *
12
from loguru import logger
13
from pydantic import BaseModel
14
from pydantic import Field
54
license_key: str = Field(..., title="The license key to replace")
55
56
63
-class CreateLicenseRequest(BaseModel):
64
- """
65
- A Pydantic model for creating a license.
57
+class TrialLicenseRequest(BaseModel):
58
+ period: Optional[int] = Field(7, title="The period of the trial license")
59
+ email: str = Field(..., title="The email of the user")
60
+ feature_name: str = Field(..., title="The feature name")
61
+ customer_name: str = Field(..., title="The customer name")
62
+ company_name: str = Field(..., title="The company name")
63
67
- Attributes:
68
- product_id (int): The product id.
69
- notes (str): The notes.
70
- new_customer (bool): Whether the customer is new.
71
- name (str): The customer name.
72
- email (str): The customer email.
73
- company_name (str): The customer company name.
74
- """
64
76
- product_id: int = Field(24355, title="The product id")
77
- notes: str = Field("Test Key", title="The notes")
78
- new_customer: bool = Field(True, title="Whether the customer is new")
79
- name: str = Field("Test Customer", title="The customer name")
80
- email: str = Field(..., title="The customer email")
81
- company_name: str = Field("Test Company", title="The customer company name")
65
+class TrialLicenseResponse(BaseModel):
66
+ license_key: str
67
+ success: bool
68
+ message: str
69
70
71
class CreateCustomerKeyResult(BaseModel):
86
87
88
class Customer(BaseModel):
102
- Id: int
103
- Name: str
104
- Email: str
105
- CompanyName: str
106
- Created: int
89
+ id: int
90
+ name: str
91
+ email: str
92
+ companyName: str
93
+ created: dt
94
95
96
class RawResponse(BaseModel):
102
103
104
class LicenseResponse(BaseModel):
118
- product_id: int
105
+ productId: int
106
id: int
107
key: str
108
created: dt
118
f8: bool
119
notes: str
120
block: bool
134
- global_id: int
121
+ globalId: int
122
customer: Customer
136
- activated_machines: List
137
- trial_activation: bool
138
- max_no_of_machines: int
139
- allowed_machines: Optional[Any]
140
- data_objects: List
141
- sign_date: dt
142
- reseller: Optional[Any]
123
+ activatedMachines: List
124
+ trialActivation: bool
125
+ maxNoOfMachines: int
126
+ allowedMachines: Optional[Any]
127
+ dataObjects: List
128
+ signDate: dt
129
+ reseller: Optional[Any] = None
130
131
132
class VerifyLicenseResponse(BaseModel):
147
message: str
148
149
163
-class Feature(Enum):
164
- MIMECAST = "MIMECAST"
165
- SAP_SIEM = "SAP SIEM"
166
- HUNTRESS = "HUNTRESS"
167
- REPORTING = "REPORTING"
168
- # Add more features as needed
169
-
170
- @classmethod
171
- def get_feature_name(cls, feature_name):
172
- feature_map = {
173
- cls.MIMECAST.value: "MIMECAST",
174
- cls.SAP_SIEM.value: "SAP SIEM",
175
- cls.HUNTRESS.value: "HUNTRESS",
176
- cls.REPORTING.value: "REPORTING",
177
- # Add more mappings as needed
178
- }
179
- return feature_map.get(feature_name)
150
+class IsFeatureEnabledResponse(BaseModel):
151
+ enabled: bool
152
+ success: bool
153
+ message: str
154
155
182
-class SubscriptionCatalog(str, Enum):
183
- """
184
- The subscription catalog.
185
- """
156
+class Feature(BaseModel):
157
+ id: int
158
+ subscription_price_id: str
159
+ name: str
160
+ price: int
161
+ currency: str
162
+ info: str
163
+ short_description: str
164
+ full_description: str
165
187
- MIMECAST = (
188
- "Integrate your SIEM stack with Mimecast to detect and respond to advanced threats."
189
- "This integration includes ingesting of Mimecast logs into your SIEM stack, Grafana dashboards,"
190
- "and alerts for advanced threat detection.",
191
- )
192
- HUNTRESS = "Integrate your SIEM stack with Huntress to detect and respond to advanced threats."
166
+
167
+class GetSubscriptionCatalogFeaturesResponse(BaseModel):
168
+ features: List[Feature]
169
+ success: bool
170
+ message: str
171
+
172
+
173
+class FeatureSubscriptionRequest(BaseModel):
174
+ feature_id: int = Field(..., example=1)
175
+ cancel_url: str = Field(..., example="https://example.com/cancel")
176
+ success_url: str = Field(..., example="https://example.com/success")
177
+ customer_email: str = Field(..., example="info@socfortress.co")
178
+ company_name: str = Field(..., example="SOCFORTRESS")
179
+
180
+
181
+class GetLicenseByEmailRequest(BaseModel):
182
+ email: str = Field(..., example="info@socfortress.co")
183
+
184
+
185
+class AddLicenseToDB(BaseModel):
186
+ customer_name: str
187
+ customer_email: str
188
+ company_name: str
189
+
190
+
191
+###### ! CREATE SESSION CHECKOUT ! ######
192
+class AutomaticTax(BaseModel):
193
+ enabled: bool
194
+ liability: Optional[str] = None
195
+ status: Optional[str] = None
196
+
197
+
198
+class CustomText(BaseModel):
199
+ after_submit: Optional[str] = None
200
+ shipping_address: Optional[str] = None
201
+ submit: Optional[str] = None
202
+ terms_of_service_acceptance: Optional[str] = None
203
+
204
+
205
+class InvoiceData(BaseModel):
206
+ account_tax_ids: Optional[str] = None
207
+ custom_fields: Optional[str] = None
208
+ description: Optional[str] = None
209
+ footer: Optional[str] = None
210
+ issuer: Optional[str] = None
211
+ metadata: Dict = {}
212
+ rendering_options: Optional[str] = None
213
+
214
+
215
+class InvoiceCreation(BaseModel):
216
+ enabled: bool
217
+ invoice_data: InvoiceData
218
+
219
+
220
+class PaymentMethodOptionsCard(BaseModel):
221
+ request_three_d_secure: str
222
+
223
+
224
+class PaymentMethodOptions(BaseModel):
225
+ card: PaymentMethodOptionsCard
226
+
227
+
228
+class PhoneNumberCollection(BaseModel):
229
+ enabled: bool
230
+
231
+
232
+class TotalDetails(BaseModel):
233
+ amount_discount: int
234
+ amount_shipping: int
235
+ amount_tax: int
236
+
237
+
238
+class CustomerDetails(BaseModel):
239
+ address: Optional[str] = None
240
+ email: Optional[str] = None
241
+ name: Optional[str] = None
242
+ phone: Optional[str] = None
243
+ tax_exempt: Optional[str] = None
244
+ tax_ids: Optional[str] = None
245
+
246
+
247
+class CheckoutSession(BaseModel):
248
+ after_expiration: Optional[str] = None
249
+ allow_promotion_codes: Optional[str] = None
250
+ amount_subtotal: int
251
+ amount_total: int
252
+ automatic_tax: AutomaticTax
253
+ billing_address_collection: Optional[str] = None
254
+ cancel_url: str
255
+ client_reference_id: Optional[str] = None
256
+ client_secret: Optional[str] = None
257
+ consent: Optional[str] = None
258
+ consent_collection: Optional[str] = None
259
+ created: int
260
+ currency: str
261
+ currency_conversion: Optional[str] = None
262
+ custom_fields: List = []
263
+ custom_text: CustomText
264
+ customer: Optional[str] = None
265
+ customer_creation: Optional[str] = None
266
+ customer_details: Optional[CustomerDetails] = None
267
+ customer_email: Optional[str] = None
268
+ expires_at: int
269
+ id: str
270
+ invoice: Optional[str] = None
271
+ invoice_creation: Optional[InvoiceCreation] = None
272
+ livemode: bool
273
+ locale: Optional[str] = None
274
+ metadata: Dict
275
+ mode: str
276
+ object: str
277
+ payment_intent: Optional[str] = None
278
+ payment_link: Optional[str] = None
279
+ payment_method_collection: str
280
+ payment_method_configuration_details: Optional[str] = None
281
+ payment_method_options: PaymentMethodOptions
282
+ payment_method_types: List[str]
283
+ payment_status: str
284
+ phone_number_collection: PhoneNumberCollection
285
+ recovered_from: Optional[str] = None
286
+ setup_intent: Optional[str] = None
287
+ shipping_address_collection: Optional[str] = None
288
+ shipping_cost: Optional[str] = None
289
+ shipping_details: Optional[str] = None
290
+ shipping_options: List = []
291
+ status: str
292
+ submit_type: Optional[str] = None
293
+ subscription: Optional[str] = None
294
+ success_url: str
295
+ total_details: TotalDetails
296
+ ui_mode: str
297
+ url: str
298
+
299
+
300
+class CheckoutSessionResponse(BaseModel):
301
+ success: bool = True
302
+ message: str = "Checkout session created successfully"
303
+ session: CheckoutSession
304
+
305
+
306
+class CancelSubscriptionRequest(BaseModel):
307
+ customer_email: str
308
+ subscription_price_id: str
309
+ feature_name: str
310
+
311
+
312
+class CancelSubscriptionResponse(BaseModel):
313
+ success: bool
314
+ message: str
315
316
317
license_router = APIRouter()
333
return auth
334
335
214
-def get_rsa_pub_key():
215
- rsa_public_key = os.getenv("RSA_PUBLIC_KEY")
216
- if not rsa_public_key:
217
- raise HTTPException(status_code=500, detail="RSA public key not found")
218
- return rsa_public_key
219
-
220
-
221
-def get_product_id():
222
- product_id = os.getenv("PRODUCT_ID")
223
- if not product_id:
224
- raise HTTPException(status_code=500, detail="Product id not found")
225
- return product_id
226
-
227
-
228
-def create_trial_key(auth, request):
229
- result, _ = Key.create_key(
230
- token=auth,
231
- product_id=request.product_id,
232
- period=7,
233
- notes=request.notes,
234
- new_customer=request.new_customer,
235
- name=request.name,
236
- email=request.email,
237
- company_name=request.company_name,
238
- )
239
- logger.info(result)
240
- result = CreateCustomerKeyResponseModel(response=[result])
241
- return result
242
-
243
-
244
-def create_key(auth, request):
245
- result, _ = Key.create_key(
246
- token=auth,
247
- product_id=request.product_id,
248
- period=365,
249
- notes=request.notes,
250
- new_customer=request.new_customer,
251
- name=request.name,
252
- email=request.email,
253
- company_name=request.company_name,
254
- )
255
- result = CreateCustomerKeyResponseModel(response=[result])
256
- return result
336
+async def add_license_to_db(session: AsyncSession, result, request: AddLicenseToDB):
337
+ """
338
+ Add a new license to the database.
339
340
+ :param session: AsyncSession object for the database session
341
+ :param result: The license key to be added
342
+ :param request: The request object containing customer details
343
+ :return: The newly added License object
344
+ """
345
259
-async def add_license_to_db(session: AsyncSession, result, request):
346
new_license = License(
261
- license_key=result.response[0].key,
262
- customer_name=request.name,
263
- customer_email=request.email,
347
+ license_key=result,
348
+ customer_name=request.customer_name,
349
+ customer_email=request.customer_email,
350
company_name=request.company_name,
351
)
352
+
353
logger.info(f"Adding new license: {new_license} to the database")
354
session.add(new_license)
355
await session.commit()
357
358
359
async def get_license(session: AsyncSession) -> License:
273
- try:
274
- result = await session.execute(select(License))
275
- license = result.scalars().first()
276
- if not license:
277
- raise HTTPException(status_code=404, detail="No license found")
278
- return license
279
- except Exception as e:
280
- logger.error(e)
281
- raise HTTPException(status_code=404, detail="No license found")
282
-
283
-
284
-def check_license(license: License):
285
- logger.info(f"Checking license: {license}")
286
- result, _ = Key.activate(
287
- token=get_auth_token(),
288
- rsa_pub_key=get_rsa_pub_key(),
289
- product_id=get_product_id(),
290
- key=license.license_key,
291
- machine_code=Helpers.GetMachineCode(v=2),
292
- )
293
- return result
294
-
360
+ """
361
+ Get the license from the database
362
296
-def extend_license(license: License, period: int):
297
- result, _ = Key.extend_license(
298
- token=get_auth_token(),
299
- product_id=get_product_id(),
300
- key=license.license_key,
301
- no_of_days=period,
302
- )
303
- logger.info(result)
304
- return result
363
+ :param session: The AsyncSession object for the database
364
+ :return: The License object
365
+ """
366
+ result = await session.execute(select(License))
367
+ license = result.scalars().first()
368
+ if license is None:
369
+ raise HTTPException(status_code=404, detail="No license found")
370
+ else:
371
+ return license
372
373
374
def is_license_expired(license: dict) -> bool:
381
Returns:
382
bool: True if the license is expired, False otherwise.
383
"""
317
- return dt.now() > license["expires"]
384
+ logger.info(f"License: {license}")
385
+ expires = dt.strptime(license["data"]["license"]["expires"], "%Y-%m-%dT%H:%M:%S.%f")
386
+ return dt.now() > expires
387
388
389
async def is_feature_enabled(feature_name: str, session: AsyncSession) -> bool:
399
bool: True if the feature is enabled, False otherwise.
400
"""
401
license = await get_license(session)
333
- license_details = LicenseResponse(**check_license(license).__dict__)
334
- for data_object in license_details.data_objects:
335
- if data_object["Name"] == feature_name and data_object["IntValue"] == 1:
402
+ result = await send_post_request("verify-license", data={"license_key": license.license_key})
403
+ for data_object in result["data"]["license"]["dataObjects"]:
404
+ if data_object["name"] == feature_name and data_object["intValue"] == 1:
405
return True
406
407
raise HTTPException(status_code=400, detail="Feature not enabled. You must purchase a license to use this feature.")
408
409
410
+async def send_get_request(endpoint: str) -> Dict[str, Any]:
411
+ """
412
+ Sends a GET request to the Shuffle service.
413
+
414
+ Args:
415
+ endpoint (str): The endpoint to send the GET request to.
416
+
417
+ Returns:
418
+ Dict[str, Any]: The response from the GET request.
419
+ """
420
+ logger.info(f"Sending GET request to {endpoint}")
421
+
422
+ try:
423
+ HEADERS = {
424
+ "x-api-key": f"{os.getenv('COPILOT_API_KEY')}",
425
+ "Content-Type": "application/json",
426
+ "module-version": "1.0",
427
+ }
428
+ response = requests.get(
429
+ f"https://license.socfortress.co/{endpoint}",
430
+ headers=HEADERS,
431
+ verify=False,
432
+ )
433
+
434
+ if response.status_code == 204:
435
+ return {
436
+ "data": None,
437
+ "success": True,
438
+ "message": "Successfully completed request with no content",
439
+ }
440
+ else:
441
+ return {
442
+ "data": response.json(),
443
+ "success": False if response.status_code >= 400 else True,
444
+ "message": "Successfully retrieved data",
445
+ }
446
+ except Exception as e:
447
+ logger.error(f"Failed to send GET request to {endpoint} with error: {e}")
448
+ raise HTTPException(
449
+ status_code=500,
450
+ detail=f"Failed to send GET request to {endpoint} with error: {e}",
451
+ )
452
+
453
+
454
+@license_router.get(
455
+ "/subscription_features",
456
+ description="Get the subscription features available",
457
+ response_model=GetSubscriptionCatalogFeaturesResponse,
458
+)
459
+async def get_subscription_catalog():
460
+ """
461
+ Get the subscription catalog. This is handled by the Middleware running in SOCFortress Infra
462
+
463
+ Returns:
464
+ dict: A dictionary containing the subscription catalog.
465
+ """
466
+ try:
467
+ results = await send_get_request("features")
468
+ return GetSubscriptionCatalogFeaturesResponse(
469
+ features=results["data"]["features"],
470
+ success=results["success"],
471
+ message=results["message"],
472
+ )
473
+ except Exception as e:
474
+ logger.error(e)
475
+ raise HTTPException(status_code=400, detail="Failed to get subscription features")
476
+
477
+
478
@license_router.post(
342
- "/create_trial_key",
343
- description="Create a trial license key",
479
+ "/retrieve_license_by_email",
480
+ description="Retrieve a license by email",
481
+ response_model=GetLicenseResponse,
482
)
345
-async def create_trial_license_key(request: CreateLicenseRequest, session: AsyncSession = Depends(get_db)):
483
+async def retrieve_license_by_email(request: GetLicenseByEmailRequest, session: AsyncSession = Depends(get_db)) -> GetLicenseResponse:
484
"""
347
- Create a trial license key.
485
+ Retrieve a license by email.
486
487
Args:
350
- request (CreateLicenseRequest): The request containing the license key to create.
488
+ request (GetLicenseRequest): The request containing the email to retrieve the license by.
489
session (AsyncSession, optional): The database session. Defaults to Depends(get_db).
490
491
Returns:
354
- LicenseVerificationResponse: A Pydantic model containing the verification status and message.
492
+ GetLicenseResponse: A Pydantic model containing the license key, success status, and message.
493
"""
356
- await check_if_license_exists(session)
357
- auth = get_auth_token()
358
- result = create_trial_key(auth, request)
359
- await add_license_to_db(session, result, request)
360
- return result
494
+ # Check if a license with the given email already exists in the database
495
+ result = await session.execute(select(License).where(License.customer_email == request.email))
496
+ existing_license = result.scalars().first()
497
+ if existing_license:
498
+ return GetLicenseResponse(
499
+ license_key=existing_license.license_key,
500
+ success=True,
501
+ message="License retrieved successfully",
502
+ )
503
+
504
+ results = await send_post_request("retrieve-license-by-email", data={"email": request.email})
505
+ logger.info(f"Results: {results}")
506
+ if results["data"]["success"] is False:
507
+ raise HTTPException(status_code=400, detail=f"Failed to retrieve license by email: {results['data']['message']}")
508
+
509
+ # Add the license to the database
510
+ await add_license_to_db(
511
+ session,
512
+ results["data"]["license"]["key"],
513
+ AddLicenseToDB(
514
+ customer_email=results["data"]["license"]["customer"]["email"],
515
+ customer_name=results["data"]["license"]["customer"]["name"],
516
+ company_name=results["data"]["license"]["customer"]["companyName"],
517
+ ),
518
+ )
519
+ return GetLicenseResponse(
520
+ license_key=results["data"]["license"]["key"],
521
+ success=results["data"]["success"],
522
+ message=results["data"]["message"],
523
+ )
524
525
526
@license_router.post(
364
- "/create_new_key",
365
- response_model=CreateCustomerKeyRouteResponse,
366
- description="Create a new license key",
527
+ "/create_checkout_session",
528
+ description="Create a checkout session",
529
+ response_model=CheckoutSessionResponse,
530
)
368
-async def create_new_license_key(request: CreateLicenseRequest, session: AsyncSession = Depends(get_db)) -> CreateCustomerKeyRouteResponse:
531
+async def create_checkout_session(request: FeatureSubscriptionRequest):
532
"""
370
- Create a new license key.
533
+ Create a checkout session.
534
535
Args:
373
- license_key (str): The license key to verify.
536
+ request (FeatureSubscriptionRequest): The request containing the feature id and user id.
537
+
538
+ Returns:
539
+ dict: A dictionary containing the checkout session.
540
+ """
541
+ results = await send_post_request(
542
+ "create-checkout-session",
543
+ data={
544
+ "feature_id": request.feature_id,
545
+ "cancel_url": request.cancel_url,
546
+ "success_url": request.success_url,
547
+ "customer_email": request.customer_email,
548
+ "company_name": request.company_name,
549
+ },
550
+ )
551
+ logger.info(f"Results: {results}")
552
+ if results["data"]["success"] is False:
553
+ raise HTTPException(status_code=400, detail=f"Failed to create checkout session: {results['data']['message']}")
554
+ return CheckoutSessionResponse(
555
+ session=results["data"]["session"],
556
+ success=results["data"]["success"],
557
+ message=results["data"]["message"],
558
+ )
559
+
560
+
561
+@license_router.post(
562
+ "/trial_license",
563
+ description="Create a trial license",
564
+ response_model=TrialLicenseResponse,
565
+)
566
+async def create_trial_license_key(request: TrialLicenseRequest, session: AsyncSession = Depends(get_db)) -> TrialLicenseResponse:
567
+ """
568
+ Create a trial license key.
569
+
570
+ Args:
571
+ request (CreateLicenseRequest): The request containing the license key to create.
572
session (AsyncSession, optional): The database session. Defaults to Depends(get_db).
573
574
Returns:
575
LicenseVerificationResponse: A Pydantic model containing the verification status and message.
576
"""
577
await check_if_license_exists(session)
380
- auth = get_auth_token()
381
- result = create_key(auth, request)
382
- logger.info(f"Result: {result}")
383
- await add_license_to_db(session, result, request)
384
- return CreateCustomerKeyRouteResponse(response=result.response, success=True, message="License created successfully")
578
+ results = await send_post_request(
579
+ "trial-license",
580
+ data={
581
+ "email": request.email,
582
+ "feature_name": request.feature_name,
583
+ "customer_name": request.customer_name,
584
+ "period": request.period,
585
+ "company_name": request.company_name,
586
+ },
587
+ )
588
+ logger.info(f"Results: {results}")
589
+ if results["data"]["success"] is False:
590
+ raise HTTPException(status_code=400, detail=f"Failed to create trial license: {results['data']['message']}")
591
+ await add_license_to_db(
592
+ session,
593
+ results["data"]["license_key"],
594
+ AddLicenseToDB(
595
+ customer_email=request.email,
596
+ customer_name=request.customer_name,
597
+ company_name=request.company_name,
598
+ ),
599
+ )
600
+ return TrialLicenseResponse(
601
+ license_key=results["data"]["license_key"],
602
+ success=results["data"]["success"],
603
+ message=results["data"]["message"],
604
+ )
605
606
607
@license_router.post(
388
- "/extend_license",
389
- description="Extend a license",
608
+ "/cancel_subscription",
609
+ description="Cancel a subscription",
610
+ response_model=CancelSubscriptionResponse,
611
)
391
-async def extend_license_key(period: int, session: AsyncSession = Depends(get_db)):
612
+async def cancel_subscription(request: CancelSubscriptionRequest) -> CancelSubscriptionResponse:
613
"""
393
- Extend a license key.
614
+ Cancel a subscription.
615
616
Args:
396
- period (int): The period to extend the license by.
397
- session (AsyncSession, optional): The database session. Defaults to Depends(get_db).
617
+ request (CancelSubscriptionRequest): The request containing the customer email, subscription price id, and feature name.
618
619
Returns:
400
- LicenseVerificationResponse: A Pydantic model containing the verification status and message.
620
+ dict: A dictionary containing the cancellation status.
621
"""
402
- try:
403
- license = await get_license(session)
404
- logger.info(f"License: {license}")
405
- extend_license(license, period)
406
- return {"message": "License extended successfully", "success": True}
407
- except Exception as e:
408
- logger.error(e)
409
- raise HTTPException(status_code=400, detail="License extension failed")
622
+ results = await send_post_request(
623
+ "cancel-subscription",
624
+ data={
625
+ "customer_email": request.customer_email,
626
+ "subscription_price_id": request.subscription_price_id,
627
+ "feature_name": request.feature_name,
628
+ },
629
+ )
630
+ logger.info(f"Results: {results}")
631
+ if results["data"]["success"] is False:
632
+ raise HTTPException(status_code=400, detail=f"Failed to cancel subscription: {results['data']['message']}")
633
+ return CancelSubscriptionResponse(
634
+ success=results["data"]["success"],
635
+ message=results["data"]["message"],
636
+ )
637
638
639
@license_router.get(
653
"""
654
license = await get_license(session)
655
try:
429
- logger.info(f"License: {license}")
430
- result = check_license(license)
431
- result = result.__dict__
432
- logger.info(result)
656
+ result = await send_post_request("verify-license", data={"license_key": license.license_key})
657
if is_license_expired(result):
658
raise HTTPException(status_code=400, detail="License is expired")
435
- return VerifyLicenseResponse(license=result, success=True, message="License verified successfully")
659
+ return VerifyLicenseResponse(license=result["data"]["license"], success=True, message="License verified successfully")
660
except Exception as e:
661
logger.error(e)
662
raise HTTPException(status_code=400, detail="License verification failed")
680
return GetLicenseResponse(license_key=license.license_key, success=True, message="License retrieved successfully")
681
682
683
+async def send_post_request(endpoint: str, data: Dict[str, Any] = None) -> Dict[str, Any]:
684
+ """
685
+ Sends a POST request to the Shuffle service.
686
+
687
+ Args:
688
+ endpoint (str): The endpoint to send the POST request to.
689
+ data (Dict[str, Any]): The data to send with the POST request.
690
+ connector_name (str, optional): The name of the connector to use. Defaults to "Shuffle".
691
+
692
+ Returns:
693
+ Dict[str, Any]: The response from the POST request.
694
+ """
695
+ logger.info(f"Sending POST request to {endpoint}")
696
+
697
+ try:
698
+ HEADERS = {
699
+ "x-api-key": f"{os.getenv('COPILOT_API_KEY')}",
700
+ "Content-Type": "application/json",
701
+ "module-version": "1.0",
702
+ }
703
+ response = requests.post(
704
+ f"https://license.socfortress.co/{endpoint}",
705
+ headers=HEADERS,
706
+ json=data,
707
+ verify=False,
708
+ )
709
+
710
+ if response.status_code == 200:
711
+ return {
712
+ "data": response.json(),
713
+ "success": True,
714
+ "message": "Successfully retrieved data",
715
+ }
716
+ else:
717
+ return {
718
+ "success": False,
719
+ "message": f"Failed to send POST request to {endpoint}",
720
+ }
721
+ except Exception as e:
722
+ logger.error(f"Failed to send POST request to {endpoint} with error: {e}")
723
+ raise HTTPException(
724
+ status_code=500,
725
+ detail=f"Failed to send POST request to {endpoint} with error: {e}",
726
+ )
727
+
728
+
729
@license_router.get(
460
- "/get_license_features",
461
- response_model=GetLicenseFeaturesResponse,
462
- description="Get license features",
730
+ "/is_feature_enabled/{feature_name}",
731
+ response_model=IsFeatureEnabledResponse,
732
+ description="Check if a feature is enabled in a license",
733
)
464
-async def get_license_features(session: AsyncSession = Depends(get_db)) -> GetLicenseFeaturesResponse:
734
+async def is_feature_enabled_route(feature_name: str, session: AsyncSession = Depends(get_db)) -> IsFeatureEnabledResponse:
735
"""
466
- Get the features enabled in a license.
736
+ Check if a feature is enabled in a license.
737
738
Args:
739
+ feature_name (str): The feature name to check.
740
session (AsyncSession, optional): The database session. Defaults to Depends(get_db).
741
742
Returns:
472
- dict: A dictionary containing the features enabled in the license.
743
+ bool: True if the feature is enabled, False otherwise.
744
"""
474
- license = await get_license(session)
475
- try:
476
- license_details = LicenseResponse(**check_license(license).__dict__)
477
- features = {}
478
- for data_object in license_details.data_objects:
479
- features[data_object["Name"]] = data_object["IntValue"]
480
- return GetLicenseFeaturesResponse(
481
- features=[feature for feature, value in features.items() if value == 1],
745
+ if await is_feature_enabled(feature_name, session):
746
+ return IsFeatureEnabledResponse(
747
+ enabled=True,
748
success=True,
483
- message="License features retrieved successfully",
749
+ message="Feature is enabled",
750
+ )
751
+ else:
752
+ return IsFeatureEnabledResponse(
753
+ enabled=False,
754
+ success=True,
755
+ message="Feature is not enabled",
756
)
485
- except Exception as e:
486
- logger.error(e)
487
- raise HTTPException(status_code=400, detail="Failed to get license features")
757
758
490
-@license_router.post(
491
- "/add_feature/{feature_name}",
492
- description="Add a feature to a license",
759
+@license_router.get(
760
+ "/get_license_features",
761
+ response_model=GetLicenseFeaturesResponse,
762
+ description="Get license features",
763
)
494
-async def add_feature_to_license(feature_name: str, session: AsyncSession = Depends(get_db)):
764
+async def get_license_features(session: AsyncSession = Depends(get_db)) -> GetLicenseFeaturesResponse:
765
"""
496
- Add a feature to a license.
766
+ Get the features enabled in a license.
767
768
Args:
499
- feature_name (str): The feature name to add.
769
session (AsyncSession, optional): The database session. Defaults to Depends(get_db).
770
771
Returns:
503
- LicenseVerificationResponse: A Pydantic model containing the verification status and message.
772
+ dict: A dictionary containing the features enabled in the license.
773
"""
505
- logger.info(f"Adding feature: {feature_name} to license")
506
- # Check if the feature name is valid
507
- feature_name = Feature.get_feature_name(feature_name)
508
- if feature_name is None:
509
- logger.error("Invalid feature name")
510
- raise HTTPException(status_code=400, detail="Invalid feature name")
774
+ license = await get_license(session)
775
try:
512
- license = await get_license(session)
513
- logger.info(f"License: {license}")
514
- result, _ = Data.add_data_object_to_key(
515
- token=get_auth_token(),
516
- product_id=get_product_id(),
517
- key=license.license_key,
518
- name=feature_name,
519
- string_value=f"[{feature_name}]",
520
- check_for_duplicates=True,
521
- int_value=1,
776
+ results = await send_post_request("license-features", data={"license_key": license.license_key})
777
+ return GetLicenseFeaturesResponse(
778
+ features=results["data"]["features"],
779
+ success=results["success"],
780
+ message=results["message"],
781
)
523
- logger.info(result)
524
- return result
782
except Exception as e:
783
logger.error(e)
527
- raise HTTPException(status_code=400, detail="Feature addition failed")
784
+ raise HTTPException(status_code=400, detail="Failed to get license features")
785
786
787
@license_router.post(
804
result = await session.execute(select(License))
805
license = result.scalars().first()
806
if not license:
550
- raise HTTPException(status_code=404, detail="No license found")
807
+ # Verify the license key
808
+ license_data = await send_post_request("verify-license", data={"license_key": request.license_key})
809
+ if is_license_expired(license_data):
810
+ raise HTTPException(status_code=400, detail="License is expired")
811
+ # Create a new License object with the data from the dictionary
812
+ license = License(
813
+ license_key=license_data["data"]["license"]["key"],
814
+ customer_name=license_data["data"]["license"]["customer"]["name"],
815
+ customer_email=license_data["data"]["license"]["customer"]["email"],
816
+ company_name=license_data["data"]["license"]["customer"]["companyName"],
817
+ )
818
+ session.add(license)
819
license.license_key = request.license_key
820
await session.commit()
821
return {"message": "License replaced successfully", "success": True}
841
842
843
async def update_connector(response: ThreatIntelRegisterResponse, session: AsyncSession):
844
+ """
845
+ When Threat Intel is purchased, add the API key to the connector.
846
+ """
847
await ConnectorServices.update_connector_by_id(
848
connector_id=10,
849
connector=UpdateConnector(
852
),
853
session=session,
854
)
584
-
585
-
586
-@license_router.post(
587
- "/register_to_threat_intel",
588
- description="Register to the SOCFortress Threat Intel Feed",
589
-)
590
-async def register_to_threat_intel(
591
- request: ThreatIntelRegisterRequest,
592
- session: AsyncSession = Depends(get_db),
593
-):
594
- """
595
- Register to the SOCFortress Threat Intel Feed.
596
-
597
- Args:
598
- request (ThreatIntelRegisterRequest): The request containing the customer name.
599
-
600
- Returns:
601
- ThreatIntelRegisterResponse: A Pydantic model containing the API key, success status, and message.
602
- """
603
- logger.info(f"Registering to the SOCFortress Threat Intel Feed: {request}")
604
- try:
605
- headers = create_headers(request)
606
- payload = create_payload(request)
607
- response = ThreatIntelRegisterResponse(
608
- **requests.post(
609
- request.registration_url,
610
- headers=headers,
611
- json=payload,
612
- ).json(),
613
- )
614
- await update_connector(response, session)
615
- return ThreatIntelRegisterResponse(
616
- api_key=response.api_key,
617
- success=response.success,
618
- message=response.message,
619
- )
620
- except Exception as e:
621
- logger.error(e)
622
- raise HTTPException(status_code=500, detail="Failed to register to the SOCFortress Threat Intel Feed")