aiwaf 0.1.9.0.4__py3-none-any.whl → 0.1.9.0.6__py3-none-any.whl

aiwaf/storage.py

@@ -1,4 +1,3 @@
-import os, csv, gzip, glob
 import numpy as np
 import pandas as pd
 from django.conf import settings
@@ -22,402 +21,209 @@ def _import_models():
         # Keep models as None if can't import
         pass
 
-# Configuration
-STORAGE_MODE = getattr(settings, "AIWAF_STORAGE_MODE", "models")  # "models" or "csv"
-CSV_DATA_DIR = getattr(settings, "AIWAF_CSV_DATA_DIR", "aiwaf_data")
-FEATURE_CSV = getattr(settings, "AIWAF_CSV_PATH", os.path.join(CSV_DATA_DIR, "access_samples.csv"))
-BLACKLIST_CSV = os.path.join(CSV_DATA_DIR, "blacklist.csv")
-EXEMPTION_CSV = os.path.join(CSV_DATA_DIR, "exemptions.csv") 
-KEYWORDS_CSV = os.path.join(CSV_DATA_DIR, "keywords.csv")
-
-CSV_HEADER = [
-    "ip","path_len","kw_hits","resp_time",
-    "status_idx","burst_count","total_404","label"
-]
-
-def ensure_csv_directory():
-    """Ensure the CSV data directory exists"""
-    if STORAGE_MODE == "csv" and not os.path.exists(CSV_DATA_DIR):
-        os.makedirs(CSV_DATA_DIR)
-
-class CsvFeatureStore:
-    @staticmethod
-    def persist_rows(rows):
-        ensure_csv_directory()
-        new_file = not os.path.exists(FEATURE_CSV)
-        with open(FEATURE_CSV, "a", newline="", encoding="utf-8") as f:
-            w = csv.writer(f)
-            if new_file:
-                w.writerow(CSV_HEADER)
-            w.writerows(rows)
-
-    @staticmethod
-    def load_matrix():
-        if not os.path.exists(FEATURE_CSV):
-            return np.empty((0,6))
-        df = pd.read_csv(
-            FEATURE_CSV,
-            names=CSV_HEADER,
-            skiprows=1,
-            engine="python",
-            on_bad_lines="skip"
-        )
-        feature_cols = CSV_HEADER[1:7]
-        df[feature_cols] = df[feature_cols].apply(pd.to_numeric, errors="coerce").fillna(0)
-        return df[feature_cols].to_numpy()
-
-class DbFeatureStore:
+class ModelFeatureStore:
     @staticmethod
     def persist_rows(rows):
+        """Persist feature data to Django models"""
         _import_models()
-        if FeatureSample is not None:
-            objs = []
-            for ip,pl,kw,rt,si,bc,t404,label in rows:
-                objs.append(FeatureSample(
-                    ip=ip, path_len=pl, kw_hits=kw,
-                    resp_time=rt, status_idx=si,
-                    burst_count=bc, total_404=t404,
-                    label=label
-                ))
-            FeatureSample.objects.bulk_create(objs, ignore_conflicts=True)
-
-    @staticmethod
-    def load_matrix():
+        if FeatureSample is None:
+            print("Warning: Django models not available, skipping feature storage")
+            return
+            
+        for row in rows:
+            try:
+                FeatureSample.objects.create(
+                    ip=row[0],
+                    path_len=int(row[1]),
+                    kw_hits=int(row[2]),
+                    resp_time=float(row[3]),
+                    status_idx=int(row[4]),
+                    burst_count=int(row[5]),
+                    total_404=int(row[6]),
+                    label=int(row[7]),
+                    created_at=timezone.now()
+                )
+            except Exception as e:
+                print(f"Error saving feature sample: {e}")
+
+    @staticmethod
+    def get_all_data():
+        """Get all feature data as DataFrame"""
         _import_models()
-        if FeatureSample is not None:
-            qs = FeatureSample.objects.all().values_list(
-                "path_len","kw_hits","resp_time","status_idx","burst_count","total_404"
+        if FeatureSample is None:
+            return pd.DataFrame()
+            
+        try:
+            queryset = FeatureSample.objects.all().values(
+                'ip', 'path_len', 'kw_hits', 'resp_time', 
+                'status_idx', 'burst_count', 'total_404', 'label'
             )
-            return np.array(list(qs), dtype=float)
-        return np.empty((0,6))
-
-def get_store():
-    if getattr(settings, "AIWAF_FEATURE_STORE", "csv") == "db":
-        return DbFeatureStore
-    return CsvFeatureStore
-
-
-# ============= CSV Storage Classes =============
+            df = pd.DataFrame(list(queryset))
+            if df.empty:
+                return df
+            
+            # Ensure proper column order and types
+            feature_cols = ['path_len', 'kw_hits', 'resp_time', 'status_idx', 'burst_count', 'total_404']
+            for col in feature_cols:
+                if col in df.columns:
+                    df[col] = pd.to_numeric(df[col], errors='coerce')
+            
+            return df
+        except Exception as e:
+            print(f"Error loading feature data: {e}")
+            return pd.DataFrame()
 
-class CsvBlacklistStore:
-    """CSV-based storage for IP blacklist entries"""
-    
-    @staticmethod
-    def add_ip(ip_address, reason):
-        ensure_csv_directory()
-        # Check if IP already exists
-        if CsvBlacklistStore.is_blocked(ip_address):
-            return
-        
-        # Add new entry
-        new_file = not os.path.exists(BLACKLIST_CSV)
-        with open(BLACKLIST_CSV, "a", newline="", encoding="utf-8") as f:
-            writer = csv.writer(f)
-            if new_file:
-                writer.writerow(["ip_address", "reason", "created_at"])
-            writer.writerow([ip_address, reason, timezone.now().isoformat()])
-    
+class ModelBlacklistStore:
     @staticmethod
-    def is_blocked(ip_address):
-        if not os.path.exists(BLACKLIST_CSV):
+    def is_blocked(ip):
+        """Check if IP is in blacklist"""
+        _import_models()
+        if BlacklistEntry is None:
+            return False
+        try:
+            return BlacklistEntry.objects.filter(ip_address=ip).exists()
+        except Exception:
             return False
-        
-        with open(BLACKLIST_CSV, "r", newline="", encoding="utf-8") as f:
-            reader = csv.DictReader(f)
-            for row in reader:
-                if row["ip_address"] == ip_address:
-                    return True
-        return False
-    
-    @staticmethod
-    def get_all():
-        """Return list of dictionaries with blacklist entries"""
-        if not os.path.exists(BLACKLIST_CSV):
-            return []
-        
-        entries = []
-        with open(BLACKLIST_CSV, "r", newline="", encoding="utf-8") as f:
-            reader = csv.DictReader(f)
-            for row in reader:
-                entries.append(row)
-        return entries
-    
-    @staticmethod
-    def remove_ip(ip_address):
-        if not os.path.exists(BLACKLIST_CSV):
-            return
-        
-        # Read all entries except the one to remove
-        entries = []
-        with open(BLACKLIST_CSV, "r", newline="", encoding="utf-8") as f:
-            reader = csv.DictReader(f)
-            entries = [row for row in reader if row["ip_address"] != ip_address]
-        
-        # Write back the filtered entries
-        with open(BLACKLIST_CSV, "w", newline="", encoding="utf-8") as f:
-            if entries:
-                writer = csv.DictWriter(f, fieldnames=["ip_address", "reason", "created_at"])
-                writer.writeheader()
-                writer.writerows(entries)
-
 
-class CsvExemptionStore:
-    """CSV-based storage for IP exemption entries"""
-    
     @staticmethod
-    def add_ip(ip_address, reason=""):
-        ensure_csv_directory()
-        
-        # Check if IP already exists to avoid duplicates
-        if CsvExemptionStore.is_exempted(ip_address):
+    def block_ip(ip, reason="Automated block"):
+        """Add IP to blacklist"""
+        _import_models()
+        if BlacklistEntry is None:
+            print(f"Warning: Cannot block IP {ip}, models not available")
             return
-        
-        # Add new entry
-        new_file = not os.path.exists(EXEMPTION_CSV)
         try:
-            with open(EXEMPTION_CSV, "a", newline="", encoding="utf-8") as f:
-                writer = csv.writer(f)
-                if new_file:
-                    writer.writerow(["ip_address", "reason", "created_at"])
-                writer.writerow([ip_address, reason, timezone.now().isoformat()])
+            BlacklistEntry.objects.get_or_create(
+                ip_address=ip,
+                defaults={'reason': reason, 'created_at': timezone.now()}
+            )
         except Exception as e:
-            print(f"Error writing to exemption CSV: {e}")
-            print(f"File path: {EXEMPTION_CSV}")
-            print(f"Directory exists: {os.path.exists(CSV_DATA_DIR)}")
-            raise
-    
+            print(f"Error blocking IP {ip}: {e}")
+
     @staticmethod
-    def is_exempted(ip_address):
-        if not os.path.exists(EXEMPTION_CSV):
-            return False
-        
+    def unblock_ip(ip):
+        """Remove IP from blacklist"""
+        _import_models()
+        if BlacklistEntry is None:
+            return
         try:
-            with open(EXEMPTION_CSV, "r", newline="", encoding="utf-8") as f:
-                reader = csv.DictReader(f)
-                for row in reader:
-                    if row["ip_address"] == ip_address:
-                        return True
+            BlacklistEntry.objects.filter(ip_address=ip).delete()
         except Exception as e:
-            print(f"Error reading exemption CSV: {e}")
-            return False
-        return False
-    
-    @staticmethod
-    def get_all():
-        """Return list of dictionaries with exemption entries"""
-        if not os.path.exists(EXEMPTION_CSV):
-            return []
-        
-        entries = []
-        with open(EXEMPTION_CSV, "r", newline="", encoding="utf-8") as f:
-            reader = csv.DictReader(f)
-            for row in reader:
-                entries.append(row)
-        return entries
-    
-    @staticmethod
-    def remove_ip(ip_address):
-        if not os.path.exists(EXEMPTION_CSV):
-            return
-        
-        # Read all entries except the one to remove
-        entries = []
-        with open(EXEMPTION_CSV, "r", newline="", encoding="utf-8") as f:
-            reader = csv.DictReader(f)
-            entries = [row for row in reader if row["ip_address"] != ip_address]
-        
-        # Write back the filtered entries
-        with open(EXEMPTION_CSV, "w", newline="", encoding="utf-8") as f:
-            if entries:
-                writer = csv.DictWriter(f, fieldnames=["ip_address", "reason", "created_at"])
-                writer.writeheader()
-                writer.writerows(entries)
-
+            print(f"Error unblocking IP {ip}: {e}")
 
-class CsvKeywordStore:
-    """CSV-based storage for dynamic keywords"""
-    
-    @staticmethod
-    def add_keyword(keyword, count=1):
-        ensure_csv_directory()
-        
-        # Read existing keywords
-        keywords = CsvKeywordStore._load_keywords()
-        
-        # Update or add keyword
-        keywords[keyword] = keywords.get(keyword, 0) + count
-        
-        # Save back to file
-        CsvKeywordStore._save_keywords(keywords)
-    
     @staticmethod
-    def get_top_keywords(limit=10):
-        keywords = CsvKeywordStore._load_keywords()
-        # Sort by count in descending order and return top N
-        sorted_keywords = sorted(keywords.items(), key=lambda x: x[1], reverse=True)
-        return [kw for kw, count in sorted_keywords[:limit]]
-    
-    @staticmethod
-    def remove_keyword(keyword):
-        keywords = CsvKeywordStore._load_keywords()
-        if keyword in keywords:
-            del keywords[keyword]
-            CsvKeywordStore._save_keywords(keywords)
-    
-    @staticmethod
-    def clear_all():
-        if os.path.exists(KEYWORDS_CSV):
-            os.remove(KEYWORDS_CSV)
-    
-    @staticmethod
-    def _load_keywords():
-        """Load keywords from CSV file as a dictionary"""
-        if not os.path.exists(KEYWORDS_CSV):
-            return {}
-        
-        keywords = {}
-        with open(KEYWORDS_CSV, "r", newline="", encoding="utf-8") as f:
-            reader = csv.DictReader(f)
-            for row in reader:
-                keywords[row["keyword"]] = int(row["count"])
-        return keywords
-    
-    @staticmethod
-    def _save_keywords(keywords):
-        """Save keywords dictionary to CSV file"""
-        with open(KEYWORDS_CSV, "w", newline="", encoding="utf-8") as f:
-            writer = csv.writer(f)
-            writer.writerow(["keyword", "count", "last_updated"])
-            for keyword, count in keywords.items():
-                writer.writerow([keyword, count, timezone.now().isoformat()])
-
-
-# ============= Storage Factory Functions =============
-
-def get_blacklist_store():
-    """Return appropriate blacklist storage class based on settings"""
-    if STORAGE_MODE == "csv":
-        return CsvBlacklistStore
-    else:
-        # Return a wrapper for Django models (only if models are available)
-        if BlacklistEntry is not None:
-            return ModelBlacklistStore
-        else:
-            # Fallback to CSV if models aren't available
-            return CsvBlacklistStore
-
-
-def get_exemption_store():
-    """Return appropriate exemption storage class based on settings"""
-    if STORAGE_MODE == "csv":
-        return CsvExemptionStore
-    else:
-        if IPExemption is not None:
-            return ModelExemptionStore
-        else:
-            return CsvExemptionStore
-
-
-def get_keyword_store():
-    """Return appropriate keyword storage class based on settings"""
-    if STORAGE_MODE == "csv":
-        return CsvKeywordStore
-    else:
-        if DynamicKeyword is not None:
-            return ModelKeywordStore
-        else:
-            return CsvKeywordStore
-
-
-# ============= Django Model Wrappers =============
-
-class ModelBlacklistStore:
-    """Django model-based storage for blacklist entries"""
-    
-    @staticmethod
-    def add_ip(ip_address, reason):
+    def get_all_blocked_ips():
+        """Get all blocked IPs"""
         _import_models()
-        if BlacklistEntry is not None:
-            BlacklistEntry.objects.get_or_create(ip_address=ip_address, defaults={"reason": reason})
-    
-    @staticmethod
-    def is_blocked(ip_address):
-        _import_models()
-        if BlacklistEntry is not None:
-            return BlacklistEntry.objects.filter(ip_address=ip_address).exists()
-        return False
-    
-    @staticmethod
-    def get_all():
-        _import_models()
-        if BlacklistEntry is not None:
-            return list(BlacklistEntry.objects.values("ip_address", "reason", "created_at"))
-        return []
-    
-    @staticmethod
-    def remove_ip(ip_address):
-        _import_models()
-        if BlacklistEntry is not None:
-            BlacklistEntry.objects.filter(ip_address=ip_address).delete()
-
+        if BlacklistEntry is None:
+            return []
+        try:
+            return list(BlacklistEntry.objects.values_list('ip_address', flat=True))
+        except Exception:
+            return []
 
 class ModelExemptionStore:
-    """Django model-based storage for exemption entries"""
-    
     @staticmethod
-    def add_ip(ip_address, reason=""):
+    def is_exempted(ip):
+        """Check if IP is exempted"""
         _import_models()
-        if IPExemption is not None:
-            IPExemption.objects.get_or_create(ip_address=ip_address, defaults={"reason": reason})
-    
+        if IPExemption is None:
+            return False
+        try:
+            return IPExemption.objects.filter(ip_address=ip).exists()
+        except Exception:
+            return False
+
     @staticmethod
-    def is_exempted(ip_address):
+    def add_exemption(ip, reason="Manual exemption"):
+        """Add IP to exemption list"""
         _import_models()
-        if IPExemption is not None:
-            return IPExemption.objects.filter(ip_address=ip_address).exists()
-        return False
-    
+        if IPExemption is None:
+            print(f"Warning: Cannot exempt IP {ip}, models not available")
+            return
+        try:
+            IPExemption.objects.get_or_create(
+                ip_address=ip,
+                defaults={'reason': reason, 'created_at': timezone.now()}
+            )
+        except Exception as e:
+            print(f"Error exempting IP {ip}: {e}")
+
     @staticmethod
-    def get_all():
+    def remove_exemption(ip):
+        """Remove IP from exemption list"""
         _import_models()
-        if IPExemption is not None:
-            return list(IPExemption.objects.values("ip_address", "reason", "created_at"))
-        return []
-    
+        if IPExemption is None:
+            return
+        try:
+            IPExemption.objects.filter(ip_address=ip).delete()
+        except Exception as e:
+            print(f"Error removing exemption for IP {ip}: {e}")
+
     @staticmethod
-    def remove_ip(ip_address):
+    def get_all_exempted_ips():
+        """Get all exempted IPs"""
         _import_models()
-        if IPExemption is not None:
-            IPExemption.objects.filter(ip_address=ip_address).delete()
-
+        if IPExemption is None:
+            return []
+        try:
+            return list(IPExemption.objects.values_list('ip_address', flat=True))
+        except Exception:
+            return []
 
 class ModelKeywordStore:
-    """Django model-based storage for dynamic keywords"""
-    
     @staticmethod
-    def add_keyword(keyword, count=1):
+    def add_keyword(keyword):
+        """Add a keyword to the dynamic keyword list"""
         _import_models()
-        if DynamicKeyword is not None:
-            obj, created = DynamicKeyword.objects.get_or_create(keyword=keyword, defaults={"count": count})
+        if DynamicKeyword is None:
+            return
+        try:
+            obj, created = DynamicKeyword.objects.get_or_create(keyword=keyword)
             if not created:
-                obj.count += count
+                obj.count += 1
                 obj.save()
-    
-    @staticmethod
-    def get_top_keywords(limit=10):
-        _import_models()
-        if DynamicKeyword is not None:
-            return list(DynamicKeyword.objects.order_by("-count").values_list("keyword", flat=True)[:limit])
-        return []
-    
+        except Exception as e:
+            print(f"Error adding keyword {keyword}: {e}")
+
     @staticmethod
-    def remove_keyword(keyword):
+    def get_top_keywords(n=10):
+        """Get top N keywords by count"""
         _import_models()
-        if DynamicKeyword is not None:
-            DynamicKeyword.objects.filter(keyword=keyword).delete()
-    
+        if DynamicKeyword is None:
+            return []
+        try:
+            return list(
+                DynamicKeyword.objects.order_by('-count')[:n]
+                .values_list('keyword', flat=True)
+            )
+        except Exception:
+            return []
+
     @staticmethod
-    def clear_all():
+    def reset_keywords():
+        """Reset all keyword counts"""
         _import_models()
-        if DynamicKeyword is not None:
+        if DynamicKeyword is None:
+            return
+        try:
             DynamicKeyword.objects.all().delete()
+        except Exception as e:
+            print(f"Error resetting keywords: {e}")
+
+# Factory functions that only return Django model stores
+def get_feature_store():
+    """Get the feature store (Django models only)"""
+    return ModelFeatureStore()
+
+def get_blacklist_store():
+    """Get the blacklist store (Django models only)"""
+    return ModelBlacklistStore()
+
+def get_exemption_store():
+    """Get the exemption store (Django models only)"""
+    return ModelExemptionStore()
+
+def get_keyword_store():
+    """Get the keyword store (Django models only)"""
+    return ModelKeywordStore()

aiwaf/trainer.py

@@ -78,18 +78,6 @@ def _read_all_logs() -> list[str]:
             except OSError:
                 continue
     
-    # If no lines found from main log, try AI-WAF middleware CSV log
-    if not lines:
-        middleware_csv = getattr(settings, "AIWAF_MIDDLEWARE_LOG", "aiwaf_requests.log").replace('.log', '.csv')
-        if os.path.exists(middleware_csv):
-            try:
-                from .middleware_logger import AIWAFCSVLogParser
-                csv_lines = AIWAFCSVLogParser.get_log_lines_for_trainer(middleware_csv)
-                lines.extend(csv_lines)
-                print(f"📋 Using AI-WAF middleware CSV log: {middleware_csv} ({len(csv_lines)} entries)")
-            except Exception as e:
-                print(f"⚠️  Failed to read middleware CSV log: {e}")
-    
     return lines