driftwatch-cli 0.1.0__py3-none-any.whl

drift_engine/aws_client.py

@@ -0,0 +1,241 @@
+import sys
+import boto3
+from datetime import datetime, timedelta
+
+def fetch_live_ec2_instances(region: str) -> dict:
+    ec2 = boto3.client("ec2", region_name=region)
+    live = {}
+    try:
+        paginator = ec2.get_paginator("describe_instances")
+        for page in paginator.paginate():
+            for reservation in page.get("Reservations", []):
+                for instance in reservation.get("Instances", []):
+                    if instance.get("State", {}).get("Name") == "terminated":
+                        continue
+                    
+                    tags_list = instance.get("Tags", [])
+                    tags_dict = {t["Key"]: t["Value"] for t in tags_list if "Key" in t and "Value" in t}
+                    name = tags_dict.get("Name", "Unknown")
+                    
+                    sg_ids = []
+                    for sg in instance.get("SecurityGroups", []):
+                        if "GroupId" in sg:
+                            sg_ids.append(sg.get("GroupId"))
+                    sg_ids.sort()
+                    
+                    instance_id = instance["InstanceId"]
+                    live[instance_id] = {
+                        "type": "aws_instance",
+                        "name": name,
+                        "attributes": {
+                            "id": instance_id,
+                            "instance_type": instance.get("InstanceType"),
+                            "ami": instance.get("ImageId"),
+                            "tags": tags_dict,
+                            "vpc_security_group_ids": sg_ids
+                        },
+                    }
+    except Exception as e:
+        print(f"Failed to fetch aws_instance: {e}", file=sys.stderr)
+        return None
+    return live
+
+def fetch_live_s3_buckets(region: str) -> dict:
+    s3 = boto3.client("s3", region_name=region)
+    live = {}
+    try:
+        response = s3.list_buckets()
+        for bucket in response.get("Buckets", []):
+            bucket_name = bucket["Name"]
+            try:
+                tags_response = s3.get_bucket_tagging(Bucket=bucket_name)
+                tags_list = tags_response.get("TagSet", [])
+                tags_dict = {t["Key"]: t["Value"] for t in tags_list if "Key" in t and "Value" in t}
+                name = tags_dict.get("Name", bucket_name)
+            except Exception:
+                tags_dict = {}
+                name = bucket_name
+
+            live[bucket_name] = {
+                "type": "aws_s3_bucket",
+                "name": name,
+                "attributes": {
+                    "id": bucket_name,
+                    "bucket": bucket_name,
+                    "tags": tags_dict,
+                },
+            }
+    except Exception as e:
+        print(f"Failed to fetch aws_s3_bucket: {e}", file=sys.stderr)
+        return None
+    return live
+
+def fetch_live_security_groups(region: str) -> dict:
+    ec2 = boto3.client("ec2", region_name=region)
+    live = {}
+    try:
+        paginator = ec2.get_paginator("describe_security_groups")
+        for page in paginator.paginate():
+            for sg in page.get("SecurityGroups", []):
+                sg_id = sg["GroupId"]
+                sg_name = sg.get("GroupName", "")
+                
+                tags_list = sg.get("Tags", [])
+                tags_dict = {t["Key"]: t["Value"] for t in tags_list if "Key" in t and "Value" in t}
+                name_tag = tags_dict.get("Name", sg_name)
+                
+                ingress_rules = []
+                for perm in sg.get("IpPermissions", []):
+                    cidrs = [ip.get("CidrIp") for ip in perm.get("IpRanges", []) if ip.get("CidrIp")]
+                    if cidrs:
+                        ingress_rules.append({
+                            "from_port": perm.get("FromPort", 0),
+                            "to_port": perm.get("ToPort", 0),
+                            "protocol": perm.get("IpProtocol", "-1"),
+                            "cidr_blocks": cidrs
+                        })
+                
+                egress_rules = []
+                for perm in sg.get("IpPermissionsEgress", []):
+                    cidrs = [ip.get("CidrIp") for ip in perm.get("IpRanges", []) if ip.get("CidrIp")]
+                    if cidrs:
+                        egress_rules.append({
+                            "from_port": perm.get("FromPort", 0),
+                            "to_port": perm.get("ToPort", 0),
+                            "protocol": perm.get("IpProtocol", "-1"),
+                            "cidr_blocks": cidrs
+                        })
+                
+                live[sg_id] = {
+                    "type": "aws_security_group",
+                    "name": name_tag,
+                    "attributes": {
+                        "id": sg_id,
+                        "name": sg_name,
+                        "description": sg.get("Description", ""),
+                        "tags": tags_dict,
+                        "ingress": ingress_rules,
+                        "egress": egress_rules,
+                    },
+                }
+    except Exception as e:
+        print(f"Failed to fetch aws_security_group: {e}", file=sys.stderr)
+        return None
+    return live
+
+def fetch_live_rds_instances(region: str) -> dict:
+    rds = boto3.client("rds", region_name=region)
+    live = {}
+    try:
+        paginator = rds.get_paginator("describe_db_instances")
+        for page in paginator.paginate():
+            for db in page.get("DBInstances", []):
+                db_id = db.get("DBInstanceIdentifier")
+                if not db_id:
+                    continue
+                live[db_id] = {
+                    "type": "aws_db_instance",
+                    "name": db_id,
+                    "attributes": {
+                        "id": db_id,
+                        "identifier": db_id,
+                        "allocated_storage": db.get("AllocatedStorage"),
+                        "engine": db.get("Engine"),
+                        "engine_version": db.get("EngineVersion"),
+                        "instance_class": db.get("DBInstanceClass"),
+                        "multi_az": db.get("MultiAZ"),
+                    },
+                }
+    except Exception as e:
+        print(f"Failed to fetch aws_db_instance: {e}", file=sys.stderr)
+        return None
+    return live
+
+def fetch_live_lambda_functions(region: str) -> dict:
+    lambda_client = boto3.client("lambda", region_name=region)
+    live = {}
+    try:
+        paginator = lambda_client.get_paginator("list_functions")
+        for page in paginator.paginate():
+            for func in page.get("Functions", []):
+                func_name = func.get("FunctionName")
+                if not func_name:
+                    continue
+                live[func_name] = {
+                    "type": "aws_lambda_function",
+                    "name": func_name,
+                    "attributes": {
+                        "id": func_name,
+                        "function_name": func_name,
+                        "runtime": func.get("Runtime"),
+                        "handler": func.get("Handler"),
+                        "memory_size": func.get("MemorySize"),
+                        "timeout": func.get("Timeout"),
+                        "role": func.get("Role"),
+                    },
+                }
+    except Exception as e:
+        print(f"Failed to fetch aws_lambda_function: {e}", file=sys.stderr)
+        return None
+    return live
+
+def fetch_live_iam_roles(region: str) -> dict:
+    iam = boto3.client("iam", region_name=region)
+    live = {}
+    try:
+        paginator = iam.get_paginator("list_roles")
+        for page in paginator.paginate():
+            for role in page.get("Roles", []):
+                role_name = role.get("RoleName")
+                if not role_name:
+                    continue
+                if role_name.startswith("AWSServiceRoleFor") or role.get("Path", "").startswith("/aws-service-role/"):
+                    continue
+                
+                try:
+                    policies = iam.list_attached_role_policies(RoleName=role_name)
+                    attached_policies = sorted(
+                        p["PolicyArn"] for p in policies.get("AttachedPolicies", []) if "PolicyArn" in p
+                    )
+                except Exception:
+                    attached_policies = []
+                
+                live[role_name] = {
+                    "type": "aws_iam_role",
+                    "name": role_name,
+                    "attributes": {
+                        "id": role_name,
+                        "name": role_name,
+                        "path": role.get("Path", "/"),
+                        "arn": role.get("Arn", ""),
+                        "attached_policies": attached_policies
+                    }
+                }
+    except Exception as e:
+        print(f"Failed to fetch aws_iam_role: {e}", file=sys.stderr)
+        return None
+    return live
+
+def get_resource_cost(resource_id: str) -> float:
+    try:
+        client = boto3.client("ce", region_name="us-east-1")
+        
+        end_date = datetime.today().strftime("%Y-%m-%d")
+        start_date = (datetime.today() - timedelta(days=30)).strftime("%Y-%m-%d")
+        
+        response = client.get_cost_and_usage(
+            TimePeriod={"Start": start_date, "End": end_date},
+            Granularity="MONTHLY",
+            Metrics=["UnblendedCost"],
+            Filter={
+                "Dimensions": {
+                    "Key": "RESOURCE_ID",
+                    "Values": [resource_id]
+                }
+            }
+        )
+        
+        usd_cost = float(response["ResultsByTime"][0]["Total"]["UnblendedCost"]["Amount"])
+        return usd_cost
+    except Exception:
+        return 0.0

drift_engine/core.py

@@ -0,0 +1,190 @@
+import os
+import sys
+from drift_engine.models import (
+    DriftResult,
+    DriftType,
+    MONITORED_ATTRIBUTES,
+    ATTRIBUTE_SEVERITY
+)
+from drift_engine.tf_parser import load_terraform_state
+from drift_engine.aws_client import (
+    fetch_live_ec2_instances,
+    fetch_live_s3_buckets,
+    fetch_live_security_groups,
+    fetch_live_rds_instances,
+    fetch_live_lambda_functions,
+    fetch_live_iam_roles
+)
+
+SEVERITY_LEVELS = {"LOW": 0, "MEDIUM": 1, "HIGH": 2, "CRITICAL": 3}
+
+def normalize_sg_rules(rules) -> list:
+    normalized = []
+    if not isinstance(rules, list):
+        return normalized
+        
+    for rule in rules:
+        if isinstance(rule, dict):
+            cidrs = rule.get('cidr_blocks') or []
+            if isinstance(cidrs, list):
+                cidrs = tuple(sorted(cidrs))
+            else:
+                cidrs = tuple()
+                
+            normalized.append({
+                'from_port': rule.get('from_port'),
+                'to_port': rule.get('to_port'),
+                'protocol': rule.get('protocol'),
+                'cidr_blocks': cidrs
+            })
+            
+    final_rules = []
+    for t in {tuple(sorted(d.items())) for d in normalized}:
+        rule_dict = dict(t)
+        rule_dict['cidr_blocks'] = list(rule_dict['cidr_blocks'])
+        final_rules.append(rule_dict)
+        
+    return final_rules
+
+def compare_attributes(tf, live, r_type) -> dict:
+    monitored = MONITORED_ATTRIBUTES.get(r_type, set())
+    diff = {}
+    
+    for key in monitored:
+        tf_val = tf.get(key)
+        live_val = live.get(key)
+        
+        if (tf_val in [None, "", [], {}]) and (live_val in [None, "", [], {}]):
+            continue
+            
+        if r_type == "aws_security_group" and key in ["ingress", "egress"]:
+            tf_norm = normalize_sg_rules(tf_val)
+            live_norm = normalize_sg_rules(live_val)
+            if tf_norm != live_norm:
+                diff[key] = {"terraform": tf_norm, "live": live_norm}
+        elif tf_val != live_val:
+            diff[key] = {"terraform": tf_val, "live": live_val}
+            
+    return diff
+
+def detect_drift(tf_state_path: str, region: str):
+    tf_resources = load_terraform_state(tf_state_path)
+    
+    if not tf_resources:
+        return [], 0
+
+    failed_types = set()
+
+    live_ec2 = fetch_live_ec2_instances(region)
+    if live_ec2 is None:
+        failed_types.add("aws_instance")
+        live_ec2 = {}
+
+    live_s3 = fetch_live_s3_buckets(region)
+    if live_s3 is None:
+        failed_types.add("aws_s3_bucket")
+        live_s3 = {}
+
+    live_sg = fetch_live_security_groups(region)
+    if live_sg is None:
+        failed_types.add("aws_security_group")
+        live_sg = {}
+
+    live_rds = fetch_live_rds_instances(region)
+    if live_rds is None:
+        failed_types.add("aws_db_instance")
+        live_rds = {}
+
+    live_lambda = fetch_live_lambda_functions(region)
+    if live_lambda is None:
+        failed_types.add("aws_lambda_function")
+        live_lambda = {}
+
+    live_iam = fetch_live_iam_roles(region)
+    if live_iam is None:
+        failed_types.add("aws_iam_role")
+        live_iam = {}
+        
+    live_resources = {**live_ec2, **live_s3, **live_sg, **live_rds, **live_lambda, **live_iam}
+    
+    results = []
+    all_ids = set(tf_resources) | set(live_resources)
+    total_scanned = len(all_ids)
+    
+    for rid in all_ids:
+        in_tf = rid in tf_resources
+        in_live = rid in live_resources
+        
+        res_name = "Unknown"
+        if in_tf:
+            res_name = tf_resources[rid]["name"]
+        elif in_live:
+            res_name = live_resources[rid]["name"]
+            
+        if in_tf and not in_live:
+            if tf_resources[rid]["type"] in failed_types:
+                continue
+            results.append(DriftResult(
+                resource_type=tf_resources[rid]["type"],
+                resource_id=rid,
+                drift_type=DriftType.MISSING,
+                resource_name=res_name,
+                tf_attributes=tf_resources[rid]["attributes"]
+            ))
+        elif in_live and not in_tf:
+            results.append(DriftResult(
+                resource_type=live_resources[rid]["type"],
+                resource_id=rid,
+                drift_type=DriftType.UNMANAGED,
+                resource_name=res_name,
+                live_attributes=live_resources[rid]["attributes"]
+            ))
+        else:
+            diff = compare_attributes(
+                tf_resources[rid]["attributes"],
+                live_resources[rid]["attributes"],
+                tf_resources[rid]["type"]
+            )
+            if diff:
+                results.append(DriftResult(
+                    resource_type=tf_resources[rid]["type"],
+                    resource_id=rid,
+                    drift_type=DriftType.MODIFIED,
+                    resource_name=res_name,
+                    tf_attributes=tf_resources[rid]["attributes"],
+                    live_attributes=live_resources[rid]["attributes"],
+                    diff=diff
+                ))
+                
+    return results, total_scanned
+
+def get_severity(r_type: str, d_type: DriftType, diff: dict = None) -> str:
+    if d_type == DriftType.MISSING or d_type == DriftType.UNMANAGED:
+        if r_type in ["aws_security_group", "aws_iam_role"]:
+            return "CRITICAL"
+        return "HIGH"
+        
+    if diff:
+        type_severity_map = ATTRIBUTE_SEVERITY.get(r_type, {})
+        highest_sev = "LOW"
+        for attr in diff.keys():
+            attr_sev = type_severity_map.get(attr, "MEDIUM")
+            if SEVERITY_LEVELS.get(attr_sev, 1) > SEVERITY_LEVELS.get(highest_sev, 0):
+                highest_sev = attr_sev
+        return highest_sev
+
+    if r_type in ["aws_security_group", "aws_iam_role"]:
+        return "CRITICAL"
+    if r_type == "aws_instance":
+        return "HIGH"
+    return "MEDIUM"
+
+if __name__ == "__main__":
+    state_file = os.environ.get("TF_STATE_PATH", "terraform/terraform.tfstate")
+    scan_region = os.environ.get("AWS_DEFAULT_REGION", "ap-south-1")
+    print(f"Executing DriftWatch Engine on {state_file} ({scan_region})...")
+    drift_items, scanned = detect_drift(state_file, scan_region)
+    print(f"Scanned {scanned} resources. Found {len(drift_items)} drift items.")
+    for item in drift_items:
+        sev = get_severity(item.resource_type, item.drift_type, item.diff)
+        print(f" - [{item.drift_type.value}] {item.resource_type} ({item.resource_id}): Severity={sev}")

drift_engine/database.py

@@ -0,0 +1,54 @@
+import os
+import json
+import psycopg2
+
+def save_drift_to_db(drift_results: list):
+    if not drift_results:
+        return
+
+    db_user = os.environ.get("DB_USER")
+    db_password = os.environ.get("DB_PASSWORD")
+    db_name = os.environ.get("DB_NAME", "driftwatch")
+    db_host = os.environ.get("DB_HOST", "localhost")
+    db_port = os.environ.get("DB_PORT", "5432")
+
+    # Security Fix: Check if credentials exist before connecting
+    if not db_user or not db_password:
+        print("⚠️ DB_USER or DB_PASSWORD not found in environment. Skipping database save.")
+        return
+        
+    try:
+        conn = psycopg2.connect(
+            dbname=db_name,
+            user=db_user,
+            password=db_password,
+            host=db_host,
+            port=db_port
+        )
+        cursor = conn.cursor()
+
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS drift_history (
+                id SERIAL PRIMARY KEY,
+                scan_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                resource_type VARCHAR(255),
+                resource_id VARCHAR(255),
+                resource_name VARCHAR(255),
+                drift_type VARCHAR(50),
+                diff_details TEXT
+            )
+        ''')
+
+        for r in drift_results:
+            diff_json = json.dumps(r.diff) if r.diff else "{}"
+            cursor.execute('''
+                INSERT INTO drift_history (resource_type, resource_id, resource_name, drift_type, diff_details)
+                VALUES (%s, %s, %s, %s, %s)
+            ''', (r.resource_type, r.resource_id, r.resource_name, r.drift_type.value, diff_json))
+
+        conn.commit()
+        cursor.close()
+        conn.close()
+        print("✅ Drift history saved to PostgreSQL database.")
+    except Exception as e:
+        print(f"❌ Database error: {e}")

drift_engine/explain.py

@@ -0,0 +1,72 @@
+import os
+import json
+import requests
+
+def get_deterministic_remediation_suggestion(resource_type: str, resource_id: str, diff_data: dict, drift_type: str) -> str:
+    safe_name = resource_id.replace("-", "_").replace(".", "_").replace("/", "_")
+    
+    if drift_type == "UNMANAGED":
+        return (
+            f"# To bring this unmanaged {resource_type} into Terraform state:\n"
+            f"terraform import {resource_type}.{safe_name} {resource_id}"
+        )
+    elif drift_type == "MISSING":
+        return (
+            f"# To recreate this missing {resource_type} defined in IaC:\n"
+            f"terraform apply -target={resource_type}.{safe_name}"
+        )
+    elif drift_type == "MODIFIED":
+        lines = [f"# To align live {resource_type} ({resource_id}) with Terraform configuration:"]
+        if diff_data:
+            for attr, vals in diff_data.items():
+                if isinstance(vals, dict) and "terraform" in vals:
+                    lines.append(f"# Set attribute '{attr}' to: {vals['terraform']}")
+        lines.append(f"terraform apply -target={resource_type}.{safe_name}")
+        return "\n".join(lines)
+    
+    return f"terraform refresh"
+
+def get_drift_explanation(resource_type: str, resource_id: str, diff_data: dict, drift_type: str) -> str:
+    api_key = os.environ.get("GROQ_API_KEY", "").strip()
+    
+    if not api_key:
+        return "AI explanation unavailable: GROQ_API_KEY not set in environment."
+
+    url = "https://api.groq.com/openai/v1/chat/completions"
+    formatted_diff = json.dumps(diff_data, indent=2)
+    
+    prompt = (
+        f"You are a strict AWS Cloud Security and Reliability expert. Analyze the following infrastructure drift:\n"
+        f"Resource Type: {resource_type}\n"
+        f"Resource ID: {resource_id}\n"
+        f"Drift Type: {drift_type}\n"
+        f"Diff Details:\n{formatted_diff}\n\n"
+        f"Provide a concise, plain-English summary (2-3 sentences max) explaining ONLY the security risks, "
+        f"compliance implications, or operational impact of this drift. "
+        f"Do NOT generate or guess CLI commands or Terraform scripts."
+    )
+
+    payload = {
+        "model": "llama-3.1-8b-instant",
+        "messages": [
+            {"role": "system", "content": "You are an AWS infrastructure and security analyst. Provide concise risk analyses only."},
+            {"role": "user", "content": prompt}
+        ],
+        "temperature": 0.2
+    }
+    
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {api_key}"
+    }
+
+    try:
+        response = requests.post(url, json=payload, headers=headers, timeout=15)
+        response.raise_for_status()
+        data = response.json()
+        return data.get("choices", [{}])[0].get("message", {}).get("content", "No risk analysis generated.")
+    
+    except requests.exceptions.RequestException as req_err:
+        return f"AI API Network Error (Groq): {req_err}"
+    except Exception as e:
+        return f"AI API Error (Groq): {str(e)}"

drift_engine/models.py

@@ -0,0 +1,65 @@
+from dataclasses import dataclass, field
+from enum import Enum
+
+class DriftType(Enum):
+    MISSING = "MISSING"
+    MODIFIED = "MODIFIED"
+    UNMANAGED = "UNMANAGED"
+
+@dataclass
+class DriftResult:
+    resource_type: str
+    resource_id: str
+    drift_type: DriftType
+    resource_name: str
+    tf_attributes: dict = field(default_factory=dict)
+    live_attributes: dict = field(default_factory=dict)
+    diff: dict = field(default_factory=dict)
+    ai_analysis: str = ""
+
+MONITORED_ATTRIBUTES = {
+    "aws_instance": {"instance_type", "ami", "tags", "vpc_security_group_ids"},
+    "aws_security_group": {"ingress", "egress", "description"},
+    "aws_s3_bucket": {"bucket", "tags"},
+    "aws_db_instance": {"instance_class", "engine", "allocated_storage", "engine_version", "multi_az"},
+    "aws_lambda_function": {"runtime", "handler", "memory_size", "timeout", "role"},
+    "aws_iam_role": {"attached_policies", "path"}
+}
+
+MONITORED_RESOURCES = list(MONITORED_ATTRIBUTES.keys())
+
+ATTRIBUTE_SEVERITY = {
+    "aws_security_group": {
+        "ingress": "CRITICAL",
+        "egress": "CRITICAL",
+        "description": "LOW"
+    },
+    "aws_instance": {
+        "instance_type": "HIGH",
+        "ami": "MEDIUM",
+        "vpc_security_group_ids": "HIGH",
+        "tags": "LOW"
+    },
+    "aws_iam_role": {
+        "attached_policies": "CRITICAL",
+        "path": "LOW"
+    },
+    "aws_db_instance": {
+        "instance_class": "HIGH",
+        "allocated_storage": "MEDIUM",
+        "engine": "HIGH",
+        "engine_version": "MEDIUM",
+        "multi_az": "HIGH"
+    },
+    "aws_s3_bucket": {
+        "bucket": "HIGH",
+        "tags": "LOW"
+    },
+    "aws_lambda_function": {
+        "runtime": "HIGH",
+        "handler": "HIGH",
+        "memory_size": "MEDIUM",
+        "timeout": "MEDIUM",
+        "role": "HIGH"
+    }
+}