Aispm feature new

.github/workflows/pr-check.yml

@@ -162,21 +162,21 @@ jobs:
         with:
           persist-credentials: false
 
-      - name: Setting up the Package Version
+      - name: Validate Package Version
         env:
           PY_VER_FILE: "pyproject.toml"
-          RELEASE_NAME: "v1.1.1"
         shell: bash
         run: |-
-          export RELEASE_VERSION=$(echo ${{ env.RELEASE_NAME }} | sed 's|^v||g')
-          if cat ${{ env.PY_VER_FILE }} | grep 'version = "RELEASE_VERSION"' ; then
-            sed -i "s|^version = \"RELEASE_VERSION\"|version = \"${RELEASE_VERSION}\"|g" ${{ env.PY_VER_FILE }}
-            cat ${file}
-          else
+          if ! grep -q '^version = ' ${{ env.PY_VER_FILE }} ; then
+            echo "Version entry not found in the ${{ env.PY_VER_FILE }} file...!"
+            exit 1
+          fi
+          VERSION=$(grep '^version = ' ${{ env.PY_VER_FILE }} | sed 's/.*version = "\(.*\)".*/\1/')
+          if [ -z "$VERSION" ]; then
             echo "Version entry format is wrong in the ${{ env.PY_VER_FILE }} file...!"
-            cat ${file}
             exit 1
           fi
+          echo "Package version: $VERSION"
 
       - name: Setup Python
         uses: actions/setup-python@v5

javelin_cli/_internal/commands.py

@@ -8,9 +8,11 @@
     UnauthorizedError,
 )
 from javelin_sdk.models import (
+    AWSConfig,
     Gateway,
     GatewayConfig,
     JavelinConfig,
+    Customer,
     Model,
     Provider,
     ProviderConfig,
@@ -20,10 +22,83 @@
     Secrets,
     Template,
     TemplateConfig,
+    AzureConfig,
 )
 from pydantic import ValidationError
 
 
+def get_javelin_client_aispm():
+    # Path to cache.json file
+    home_dir = Path.home()
+    json_file_path = home_dir / ".javelin" / "cache.json"
+
+    # Load cache.json
+    if not json_file_path.exists():
+        raise FileNotFoundError(f"Configuration file not found: {json_file_path}")
+
+    with open(json_file_path, "r") as json_file:
+        cache_data = json.load(json_file)
+
+    # Retrieve the list of gateways
+    gateways = (
+        cache_data.get("memberships", {})
+        .get("data", [{}])[0]
+        .get("organization", {})
+        .get("public_metadata", {})
+        .get("Gateways", [])
+    )
+    if not gateways:
+        raise ValueError("No gateways found in the configuration.")
+
+    # Automatically select the first gateway (index 0)
+    selected_gateway = gateways[0]
+    base_url = selected_gateway["base_url"]
+
+    # Get organization metadata (where account_id might be stored)
+    organization = (
+        cache_data.get("memberships", {}).get("data", [{}])[0].get("organization", {})
+    )
+    org_metadata = organization.get("public_metadata", {})
+
+    # Get account_id from multiple possible locations (in order of preference):
+    # 1. Gateway's account_id field
+    # 2. Organization's public_metadata account_id
+    # 3. Extract from role_arn if provided
+    account_id = selected_gateway.get("account_id")
+    if not account_id:
+        account_id = org_metadata.get("account_id")
+
+    role_arn = selected_gateway.get("role_arn")
+
+    # Extract account_id from role ARN if still not found
+    # Format: arn:aws:iam::ACCOUNT_ID:role/ROLE_NAME
+    if role_arn and not account_id:
+        try:
+            parts = role_arn.split(":")
+            if len(parts) >= 5 and parts[2] == "iam":
+                account_id = parts[4]
+        except (IndexError, AttributeError):
+            pass
+
+    javelin_api_key = selected_gateway.get("api_key_value", "placeholder")
+
+    # Initialize and return the JavelinClient
+    config = JavelinConfig(
+        base_url=base_url,
+        javelin_api_key=javelin_api_key,
+    )
+
+    client = JavelinClient(config)
+
+    # Store account_id in client for AISPM service to use
+    if account_id:
+        client._aispm_account_id = account_id
+        client._aispm_user = "test-user"
+        client._aispm_userrole = "org:superadmin"
+
+    return client
+
+
 def get_javelin_client():
     # Path to cache.json file
     home_dir = Path.home()
@@ -86,6 +161,133 @@ def get_javelin_client():
     return JavelinClient(config)
 
 
+def create_customer(args):
+    client = get_javelin_client_aispm()
+    customer = Customer(
+        name=args.name,
+        description=args.description,
+        metrics_interval=args.metrics_interval,
+        security_interval=args.security_interval,
+    )
+    return client.aispm.create_customer(customer)
+
+
+def get_customer(args):
+    """
+    Gets customer details using the AISPM service.
+    """
+    try:
+        client = get_javelin_client_aispm()
+        response = client.aispm.get_customer()
+
+        # Pretty print the response for CLI output
+        formatted_response = {
+            "name": response.name,
+            "description": response.description,
+            "metrics_interval": response.metrics_interval,
+            "security_interval": response.security_interval,
+            "status": response.status,
+            "created_at": response.created_at.isoformat(),
+            "modified_at": response.modified_at.isoformat(),
+        }
+
+        print(json.dumps(formatted_response, indent=2))
+    except Exception as e:
+        print(f"Error getting customer: {e}")
+
+
+def configure_aws(args):
+    try:
+        client = get_javelin_client_aispm()
+        config = json.loads(args.config)
+        configs = [AWSConfig(**config)]
+        client.aispm.configure_aws(configs)
+        print("AWS configuration created successfully.")
+    except Exception as e:
+        print(f"Error configuring AWS: {e}")
+
+
+def get_aws_config(args):
+    """
+    Gets AWS configurations using the AISPM service.
+    """
+    try:
+        client = get_javelin_client_aispm()
+        response = client.aispm.get_aws_configs()
+        # Simply print the JSON response
+        print(json.dumps(response, indent=2))
+
+    except Exception as e:
+        print(f"Error getting AWS configurations: {e}")
+
+
+# Add these functions to commands.py
+
+
+def delete_aws_config(args):
+    """
+    Deletes an AWS configuration.
+    """
+    try:
+        client = get_javelin_client_aispm()
+        client.aispm.delete_aws_config(args.name)
+        print(f"AWS configuration '{args.name}' deleted successfully.")
+    except Exception as e:
+        print(f"Error deleting AWS config: {e}")
+
+
+def get_azure_config(args):
+    """
+    Gets Azure configurations using the AISPM service.
+    """
+    try:
+        client = get_javelin_client_aispm()
+        response = client.aispm.get_azure_config()
+        # Format and print the response nicely
+        print(json.dumps(response, indent=2))
+    except Exception as e:
+        print(f"Error getting Azure config: {e}")
+
+
+def configure_azure(args):
+    try:
+        client = get_javelin_client_aispm()
+        config = json.loads(args.config)
+        configs = [AzureConfig(**config)]
+        client.aispm.configure_azure(configs)
+        print("Azure configuration created successfully.")
+    except Exception as e:
+        print(f"Error configuring Azure: {e}")
+
+
+def get_usage(args):
+    try:
+        client = get_javelin_client_aispm()
+        usage = client.aispm.get_usage(
+            provider=args.provider,
+            cloud_account=args.account,
+            model=args.model,
+            region=args.region,
+        )
+        print(json.dumps(usage.dict(), indent=2))
+    except Exception as e:
+        print(f"Error getting usage: {e}")
+
+
+def get_alerts(args):
+    try:
+        client = get_javelin_client_aispm()
+        alerts = client.aispm.get_alerts(
+            provider=args.provider,
+            cloud_account=args.account,
+            model=args.model,
+            region=args.region,
+        )
+        print(json.dumps(alerts.dict(), indent=2))
+    except Exception as e:
+        print(f"Error getting alerts: {e}")
+
+
 def create_gateway(args):
     try:
         client = get_javelin_client()

javelin_cli/cli.py

@@ -37,11 +37,20 @@
     update_route,
     update_secret,
     update_template,
+    create_customer,
+    get_customer,
+    configure_aws,
+    configure_azure,
+    get_usage,
+    get_alerts,
+    get_aws_config,
+    get_azure_config,
+    delete_aws_config
 )
 
 
 def check_permissions():
-    """Check if user has superadmin permissions"""
+    """Check if user has permissions"""
     home_dir = Path.home()
     cache_file = home_dir / ".javelin" / "cache.json"
 
@@ -95,6 +104,83 @@ def main():
         help="Force re-authentication, overriding existing credentials",
     )
     auth_parser.set_defaults(func=authenticate)
+
+    # AISPM commands
+    aispm_parser = subparsers.add_parser("aispm", help="Manage AISPM functionality")
+    aispm_subparsers = aispm_parser.add_subparsers()
+
+    # Customer commands
+    customer_parser = aispm_subparsers.add_parser("customer", help="Manage customers")
+    customer_subparsers = customer_parser.add_subparsers()
+
+    customer_create = customer_subparsers.add_parser("create", help="Create customer")
+    customer_create.add_argument("--name", required=True, help="Customer name")
+    customer_create.add_argument("--description", help="Customer description")
+    customer_create.add_argument(
+        "--metrics-interval", default="5m", help="Metrics interval"
+    )
+    customer_create.add_argument(
+        "--security-interval", default="1m", help="Security interval"
+    )
+    customer_create.set_defaults(func=create_customer)
+
+    customer_get = customer_subparsers.add_parser("get", help="Get customer details")
+    customer_get.set_defaults(func=get_customer)
+
+    # Cloud config commands
+    config_parser = aispm_subparsers.add_parser(
+        "config", help="Manage cloud configurations"
+    )
+    config_subparsers = config_parser.add_subparsers()
+
+    aws_parser = config_subparsers.add_parser("aws", help="Configure AWS")
+    azure_parser = config_subparsers.add_parser("azure", help="Configure Azure")
+
+    aws_subparsers = aws_parser.add_subparsers()
+    aws_get_parser = aws_subparsers.add_parser("get", help="Get AWS configuration")
+    aws_get_parser.set_defaults(func=get_aws_config)
+
+    aws_create_parser = aws_subparsers.add_parser("create", help="Configure AWS")
+    aws_create_parser.add_argument(
+        "--config", type=str, required=True, help="AWS config JSON"
+    )
+    aws_create_parser.set_defaults(func=configure_aws)
+
+    aws_delete_parser = aws_subparsers.add_parser(
+        "delete", help="Delete AWS configuration"
+    )
+    aws_delete_parser.add_argument(
+        "--name", type=str, required=True, help="Name of AWS configuration to delete"
+    )
+    aws_delete_parser.set_defaults(func=delete_aws_config)
+
+    azure_subparsers = azure_parser.add_subparsers(dest="azure_command")
+    azure_get_parser = azure_subparsers.add_parser(
+        "get", help="Get Azure configuration"
+    )
+    azure_get_parser.set_defaults(func=get_azure_config)
+
+    azure_create_parser = azure_subparsers.add_parser("create", help="Configure Azure")
+    azure_create_parser.add_argument(
+        "--config", type=str, required=True, help="Azure config JSON"
+    )
+    azure_create_parser.set_defaults(func=configure_azure)
+
+    # Usage metrics
+    usage_parser = aispm_subparsers.add_parser("usage", help="Get usage metrics")
+    usage_parser.add_argument("--provider", help="Cloud provider")
+    usage_parser.add_argument("--account", help="Cloud account name")
+    usage_parser.add_argument("--model", help="Model ID")
+    usage_parser.add_argument("--region", help="Region")
+    usage_parser.set_defaults(func=get_usage)
+
+    # Alerts
+    alerts_parser = aispm_subparsers.add_parser("alerts", help="Get alerts")
+    alerts_parser.add_argument("--provider", help="Cloud provider")
+    alerts_parser.add_argument("--account", help="Cloud account name")
+    alerts_parser.add_argument("--model", help="Model ID")
+    alerts_parser.add_argument("--region", help="Region")
+    alerts_parser.set_defaults(func=get_alerts)
     # Gateway CRUD
     gateway_parser = subparsers.add_parser(
         "gateway",
@@ -415,8 +501,7 @@ def authenticate(args):
         print("Use --force to re-authenticate and override existing cache.")
         return
 
-    default_url = "https://dev.javelin.live/"
-
+    default_url = "https://dev.highflame.dev/"
     print("   O")
     print("  /|\\")
     print("  / \\    ========> Welcome to Javelin! 🚀")