aiwaf 0.1.8.6__py3-none-any.whl → 0.1.8.7__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of aiwaf might be problematic. Click here for more details.

@@ -1,14 +1,24 @@
1
- from .models import BlacklistEntry
1
+ # aiwaf/blacklist_manager.py
2
+
3
+ from .storage import get_blacklist_store
2
4
 
3
5
  class BlacklistManager:
4
6
  @staticmethod
5
7
  def block(ip, reason):
6
- BlacklistEntry.objects.get_or_create(ip_address=ip, defaults={"reason": reason})
8
+ store = get_blacklist_store()
9
+ store.add_ip(ip, reason)
7
10
 
8
11
  @staticmethod
9
12
  def is_blocked(ip):
10
- return BlacklistEntry.objects.filter(ip_address=ip).exists()
13
+ store = get_blacklist_store()
14
+ return store.is_blocked(ip)
11
15
 
12
16
  @staticmethod
13
17
  def all_blocked():
14
- return BlacklistEntry.objects.all()
18
+ store = get_blacklist_store()
19
+ return store.get_all()
20
+
21
+ @staticmethod
22
+ def unblock(ip):
23
+ store = get_blacklist_store()
24
+ store.remove_ip(ip)
@@ -1,5 +1,5 @@
1
1
  from django.core.management.base import BaseCommand, CommandError
2
- from aiwaf.models import IPExemption
2
+ from aiwaf.storage import get_exemption_store
3
3
 
4
4
  class Command(BaseCommand):
5
5
  help = 'Add an IP address to the IPExemption list (prevents blacklisting)'
@@ -11,12 +11,13 @@ class Command(BaseCommand):
11
11
  def handle(self, *args, **options):
12
12
  ip = options['ip']
13
13
  reason = options['reason']
14
- obj, created = IPExemption.objects.get_or_create(ip_address=ip, defaults={'reason': reason})
15
- if not created:
14
+
15
+ store = get_exemption_store()
16
+
17
+ if store.is_exempted(ip):
16
18
  self.stdout.write(self.style.WARNING(f'IP {ip} is already exempted.'))
17
19
  else:
20
+ store.add_ip(ip, reason)
18
21
  self.stdout.write(self.style.SUCCESS(f'IP {ip} added to exemption list.'))
19
- if reason:
20
- obj.reason = reason
21
- obj.save()
22
- self.stdout.write(self.style.SUCCESS(f'Reason set to: {reason}'))
22
+ if reason:
23
+ self.stdout.write(self.style.SUCCESS(f'Reason: {reason}'))
@@ -1,5 +1,5 @@
1
1
  from django.core.management.base import BaseCommand
2
- from aiwaf.models import BlacklistEntry, IPExemption
2
+ from aiwaf.storage import get_blacklist_store, get_exemption_store
3
3
 
4
4
  class Command(BaseCommand):
5
5
  help = 'Reset AI-WAF by clearing all blacklist and exemption (whitelist) entries'
@@ -26,9 +26,12 @@ class Command(BaseCommand):
26
26
  exemptions_only = options['exemptions_only']
27
27
  confirm = options['confirm']
28
28
 
29
+ blacklist_store = get_blacklist_store()
30
+ exemption_store = get_exemption_store()
31
+
29
32
  # Count current entries
30
- blacklist_count = BlacklistEntry.objects.count()
31
- exemption_count = IPExemption.objects.count()
33
+ blacklist_count = len(blacklist_store.get_all())
34
+ exemption_count = len(exemption_store.get_all())
32
35
 
33
36
  if blacklist_only and exemptions_only:
34
37
  self.stdout.write(self.style.ERROR('Cannot use both --blacklist-only and --exemptions-only flags'))
@@ -61,10 +64,18 @@ class Command(BaseCommand):
61
64
  deleted_counts = {'blacklist': 0, 'exemptions': 0}
62
65
 
63
66
  if clear_blacklist:
64
- deleted_counts['blacklist'], _ = BlacklistEntry.objects.all().delete()
67
+ # Clear blacklist entries
68
+ blacklist_entries = blacklist_store.get_all()
69
+ for entry in blacklist_entries:
70
+ blacklist_store.remove_ip(entry['ip_address'])
71
+ deleted_counts['blacklist'] = len(blacklist_entries)
65
72
 
66
73
  if clear_exemptions:
67
- deleted_counts['exemptions'], _ = IPExemption.objects.all().delete()
74
+ # Clear exemption entries
75
+ exemption_entries = exemption_store.get_all()
76
+ for entry in exemption_entries:
77
+ exemption_store.remove_ip(entry['ip_address'])
78
+ deleted_counts['exemptions'] = len(exemption_entries)
68
79
 
69
80
  # Report results
70
81
  if clear_blacklist and clear_exemptions:
aiwaf/middleware.py CHANGED
@@ -16,8 +16,9 @@ from django.apps import apps
16
16
  from django.urls import get_resolver
17
17
  from .trainer import STATIC_KW, STATUS_IDX, path_exists_in_django
18
18
  from .blacklist_manager import BlacklistManager
19
- from .models import DynamicKeyword, IPExemption
19
+ from .models import IPExemption
20
20
  from .utils import is_exempt, get_ip, is_ip_exempted
21
+ from .storage import get_keyword_store
21
22
 
22
23
  MODEL_PATH = getattr(
23
24
  settings,
@@ -74,15 +75,14 @@ class IPAndKeywordBlockMiddleware:
74
75
  return self.get_response(request)
75
76
  if BlacklistManager.is_blocked(ip):
76
77
  return JsonResponse({"error": "blocked"}, status=403)
78
+
79
+ keyword_store = get_keyword_store()
77
80
  segments = [seg for seg in re.split(r"\W+", path) if len(seg) > 3]
81
+
78
82
  for seg in segments:
79
- obj, _ = DynamicKeyword.objects.get_or_create(keyword=seg)
80
- DynamicKeyword.objects.filter(pk=obj.pk).update(count=F("count") + 1)
81
- dynamic_top = list(
82
- DynamicKeyword.objects
83
- .order_by("-count")
84
- .values_list("keyword", flat=True)[: getattr(settings, "AIWAF_DYNAMIC_TOP_N", 10)]
85
- )
83
+ keyword_store.add_keyword(seg)
84
+
85
+ dynamic_top = keyword_store.get_top_keywords(getattr(settings, "AIWAF_DYNAMIC_TOP_N", 10))
86
86
  all_kw = set(STATIC_KW) | set(dynamic_top)
87
87
  suspicious_kw = {
88
88
  kw for kw in all_kw
@@ -172,10 +172,11 @@ class AIAnomalyMiddleware(MiddlewareMixin):
172
172
  data.append((now, request.path, response.status_code, resp_time))
173
173
  data = [d for d in data if now - d[0] < self.WINDOW]
174
174
  cache.set(key, data, timeout=self.WINDOW)
175
+
176
+ keyword_store = get_keyword_store()
175
177
  for seg in re.split(r"\W+", request.path.lower()):
176
178
  if len(seg) > 3:
177
- obj, _ = DynamicKeyword.objects.get_or_create(keyword=seg)
178
- DynamicKeyword.objects.filter(pk=obj.pk).update(count=F("count") + 1)
179
+ keyword_store.add_keyword(seg)
179
180
 
180
181
  return response
181
182
 
aiwaf/storage.py CHANGED
@@ -2,19 +2,33 @@ import os, csv, gzip, glob
2
2
  import numpy as np
3
3
  import pandas as pd
4
4
  from django.conf import settings
5
- from .models import FeatureSample
5
+ from django.utils import timezone
6
+ from .models import FeatureSample, BlacklistEntry, IPExemption, DynamicKeyword
7
+
8
+ # Configuration
9
+ STORAGE_MODE = getattr(settings, "AIWAF_STORAGE_MODE", "models") # "models" or "csv"
10
+ CSV_DATA_DIR = getattr(settings, "AIWAF_CSV_DATA_DIR", "aiwaf_data")
11
+ FEATURE_CSV = getattr(settings, "AIWAF_CSV_PATH", os.path.join(CSV_DATA_DIR, "access_samples.csv"))
12
+ BLACKLIST_CSV = os.path.join(CSV_DATA_DIR, "blacklist.csv")
13
+ EXEMPTION_CSV = os.path.join(CSV_DATA_DIR, "exemptions.csv")
14
+ KEYWORDS_CSV = os.path.join(CSV_DATA_DIR, "keywords.csv")
6
15
 
7
- DATA_FILE = getattr(settings, "AIWAF_CSV_PATH", "access_samples.csv")
8
16
  CSV_HEADER = [
9
17
  "ip","path_len","kw_hits","resp_time",
10
18
  "status_idx","burst_count","total_404","label"
11
19
  ]
12
20
 
21
+ def ensure_csv_directory():
22
+ """Ensure the CSV data directory exists"""
23
+ if STORAGE_MODE == "csv" and not os.path.exists(CSV_DATA_DIR):
24
+ os.makedirs(CSV_DATA_DIR)
25
+
13
26
  class CsvFeatureStore:
14
27
  @staticmethod
15
28
  def persist_rows(rows):
16
- new_file = not os.path.exists(DATA_FILE)
17
- with open(DATA_FILE, "a", newline="", encoding="utf-8") as f:
29
+ ensure_csv_directory()
30
+ new_file = not os.path.exists(FEATURE_CSV)
31
+ with open(FEATURE_CSV, "a", newline="", encoding="utf-8") as f:
18
32
  w = csv.writer(f)
19
33
  if new_file:
20
34
  w.writerow(CSV_HEADER)
@@ -22,10 +36,10 @@ class CsvFeatureStore:
22
36
 
23
37
  @staticmethod
24
38
  def load_matrix():
25
- if not os.path.exists(DATA_FILE):
39
+ if not os.path.exists(FEATURE_CSV):
26
40
  return np.empty((0,6))
27
41
  df = pd.read_csv(
28
- DATA_FILE,
42
+ FEATURE_CSV,
29
43
  names=CSV_HEADER,
30
44
  skiprows=1,
31
45
  engine="python",
@@ -59,3 +73,279 @@ def get_store():
59
73
  if getattr(settings, "AIWAF_FEATURE_STORE", "csv") == "db":
60
74
  return DbFeatureStore
61
75
  return CsvFeatureStore
76
+
77
+
78
+ # ============= CSV Storage Classes =============
79
+
80
+ class CsvBlacklistStore:
81
+ """CSV-based storage for IP blacklist entries"""
82
+
83
+ @staticmethod
84
+ def add_ip(ip_address, reason):
85
+ ensure_csv_directory()
86
+ # Check if IP already exists
87
+ if CsvBlacklistStore.is_blocked(ip_address):
88
+ return
89
+
90
+ # Add new entry
91
+ new_file = not os.path.exists(BLACKLIST_CSV)
92
+ with open(BLACKLIST_CSV, "a", newline="", encoding="utf-8") as f:
93
+ writer = csv.writer(f)
94
+ if new_file:
95
+ writer.writerow(["ip_address", "reason", "created_at"])
96
+ writer.writerow([ip_address, reason, timezone.now().isoformat()])
97
+
98
+ @staticmethod
99
+ def is_blocked(ip_address):
100
+ if not os.path.exists(BLACKLIST_CSV):
101
+ return False
102
+
103
+ with open(BLACKLIST_CSV, "r", newline="", encoding="utf-8") as f:
104
+ reader = csv.DictReader(f)
105
+ for row in reader:
106
+ if row["ip_address"] == ip_address:
107
+ return True
108
+ return False
109
+
110
+ @staticmethod
111
+ def get_all():
112
+ """Return list of dictionaries with blacklist entries"""
113
+ if not os.path.exists(BLACKLIST_CSV):
114
+ return []
115
+
116
+ entries = []
117
+ with open(BLACKLIST_CSV, "r", newline="", encoding="utf-8") as f:
118
+ reader = csv.DictReader(f)
119
+ for row in reader:
120
+ entries.append(row)
121
+ return entries
122
+
123
+ @staticmethod
124
+ def remove_ip(ip_address):
125
+ if not os.path.exists(BLACKLIST_CSV):
126
+ return
127
+
128
+ # Read all entries except the one to remove
129
+ entries = []
130
+ with open(BLACKLIST_CSV, "r", newline="", encoding="utf-8") as f:
131
+ reader = csv.DictReader(f)
132
+ entries = [row for row in reader if row["ip_address"] != ip_address]
133
+
134
+ # Write back the filtered entries
135
+ with open(BLACKLIST_CSV, "w", newline="", encoding="utf-8") as f:
136
+ if entries:
137
+ writer = csv.DictWriter(f, fieldnames=["ip_address", "reason", "created_at"])
138
+ writer.writeheader()
139
+ writer.writerows(entries)
140
+
141
+
142
+ class CsvExemptionStore:
143
+ """CSV-based storage for IP exemption entries"""
144
+
145
+ @staticmethod
146
+ def add_ip(ip_address, reason=""):
147
+ ensure_csv_directory()
148
+ # Check if IP already exists
149
+ if CsvExemptionStore.is_exempted(ip_address):
150
+ return
151
+
152
+ # Add new entry
153
+ new_file = not os.path.exists(EXEMPTION_CSV)
154
+ with open(EXEMPTION_CSV, "a", newline="", encoding="utf-8") as f:
155
+ writer = csv.writer(f)
156
+ if new_file:
157
+ writer.writerow(["ip_address", "reason", "created_at"])
158
+ writer.writerow([ip_address, reason, timezone.now().isoformat()])
159
+
160
+ @staticmethod
161
+ def is_exempted(ip_address):
162
+ if not os.path.exists(EXEMPTION_CSV):
163
+ return False
164
+
165
+ with open(EXEMPTION_CSV, "r", newline="", encoding="utf-8") as f:
166
+ reader = csv.DictReader(f)
167
+ for row in reader:
168
+ if row["ip_address"] == ip_address:
169
+ return True
170
+ return False
171
+
172
+ @staticmethod
173
+ def get_all():
174
+ """Return list of dictionaries with exemption entries"""
175
+ if not os.path.exists(EXEMPTION_CSV):
176
+ return []
177
+
178
+ entries = []
179
+ with open(EXEMPTION_CSV, "r", newline="", encoding="utf-8") as f:
180
+ reader = csv.DictReader(f)
181
+ for row in reader:
182
+ entries.append(row)
183
+ return entries
184
+
185
+ @staticmethod
186
+ def remove_ip(ip_address):
187
+ if not os.path.exists(EXEMPTION_CSV):
188
+ return
189
+
190
+ # Read all entries except the one to remove
191
+ entries = []
192
+ with open(EXEMPTION_CSV, "r", newline="", encoding="utf-8") as f:
193
+ reader = csv.DictReader(f)
194
+ entries = [row for row in reader if row["ip_address"] != ip_address]
195
+
196
+ # Write back the filtered entries
197
+ with open(EXEMPTION_CSV, "w", newline="", encoding="utf-8") as f:
198
+ if entries:
199
+ writer = csv.DictWriter(f, fieldnames=["ip_address", "reason", "created_at"])
200
+ writer.writeheader()
201
+ writer.writerows(entries)
202
+
203
+
204
+ class CsvKeywordStore:
205
+ """CSV-based storage for dynamic keywords"""
206
+
207
+ @staticmethod
208
+ def add_keyword(keyword, count=1):
209
+ ensure_csv_directory()
210
+
211
+ # Read existing keywords
212
+ keywords = CsvKeywordStore._load_keywords()
213
+
214
+ # Update or add keyword
215
+ keywords[keyword] = keywords.get(keyword, 0) + count
216
+
217
+ # Save back to file
218
+ CsvKeywordStore._save_keywords(keywords)
219
+
220
+ @staticmethod
221
+ def get_top_keywords(limit=10):
222
+ keywords = CsvKeywordStore._load_keywords()
223
+ # Sort by count in descending order and return top N
224
+ sorted_keywords = sorted(keywords.items(), key=lambda x: x[1], reverse=True)
225
+ return [kw for kw, count in sorted_keywords[:limit]]
226
+
227
+ @staticmethod
228
+ def remove_keyword(keyword):
229
+ keywords = CsvKeywordStore._load_keywords()
230
+ if keyword in keywords:
231
+ del keywords[keyword]
232
+ CsvKeywordStore._save_keywords(keywords)
233
+
234
+ @staticmethod
235
+ def clear_all():
236
+ if os.path.exists(KEYWORDS_CSV):
237
+ os.remove(KEYWORDS_CSV)
238
+
239
+ @staticmethod
240
+ def _load_keywords():
241
+ """Load keywords from CSV file as a dictionary"""
242
+ if not os.path.exists(KEYWORDS_CSV):
243
+ return {}
244
+
245
+ keywords = {}
246
+ with open(KEYWORDS_CSV, "r", newline="", encoding="utf-8") as f:
247
+ reader = csv.DictReader(f)
248
+ for row in reader:
249
+ keywords[row["keyword"]] = int(row["count"])
250
+ return keywords
251
+
252
+ @staticmethod
253
+ def _save_keywords(keywords):
254
+ """Save keywords dictionary to CSV file"""
255
+ with open(KEYWORDS_CSV, "w", newline="", encoding="utf-8") as f:
256
+ writer = csv.writer(f)
257
+ writer.writerow(["keyword", "count", "last_updated"])
258
+ for keyword, count in keywords.items():
259
+ writer.writerow([keyword, count, timezone.now().isoformat()])
260
+
261
+
262
+ # ============= Storage Factory Functions =============
263
+
264
+ def get_blacklist_store():
265
+ """Return appropriate blacklist storage class based on settings"""
266
+ if STORAGE_MODE == "csv":
267
+ return CsvBlacklistStore
268
+ else:
269
+ # Return a wrapper for Django models
270
+ return ModelBlacklistStore
271
+
272
+
273
+ def get_exemption_store():
274
+ """Return appropriate exemption storage class based on settings"""
275
+ if STORAGE_MODE == "csv":
276
+ return CsvExemptionStore
277
+ else:
278
+ return ModelExemptionStore
279
+
280
+
281
+ def get_keyword_store():
282
+ """Return appropriate keyword storage class based on settings"""
283
+ if STORAGE_MODE == "csv":
284
+ return CsvKeywordStore
285
+ else:
286
+ return ModelKeywordStore
287
+
288
+
289
+ # ============= Django Model Wrappers =============
290
+
291
+ class ModelBlacklistStore:
292
+ """Django model-based storage for blacklist entries"""
293
+
294
+ @staticmethod
295
+ def add_ip(ip_address, reason):
296
+ BlacklistEntry.objects.get_or_create(ip_address=ip_address, defaults={"reason": reason})
297
+
298
+ @staticmethod
299
+ def is_blocked(ip_address):
300
+ return BlacklistEntry.objects.filter(ip_address=ip_address).exists()
301
+
302
+ @staticmethod
303
+ def get_all():
304
+ return list(BlacklistEntry.objects.values("ip_address", "reason", "created_at"))
305
+
306
+ @staticmethod
307
+ def remove_ip(ip_address):
308
+ BlacklistEntry.objects.filter(ip_address=ip_address).delete()
309
+
310
+
311
+ class ModelExemptionStore:
312
+ """Django model-based storage for exemption entries"""
313
+
314
+ @staticmethod
315
+ def add_ip(ip_address, reason=""):
316
+ IPExemption.objects.get_or_create(ip_address=ip_address, defaults={"reason": reason})
317
+
318
+ @staticmethod
319
+ def is_exempted(ip_address):
320
+ return IPExemption.objects.filter(ip_address=ip_address).exists()
321
+
322
+ @staticmethod
323
+ def get_all():
324
+ return list(IPExemption.objects.values("ip_address", "reason", "created_at"))
325
+
326
+ @staticmethod
327
+ def remove_ip(ip_address):
328
+ IPExemption.objects.filter(ip_address=ip_address).delete()
329
+
330
+
331
+ class ModelKeywordStore:
332
+ """Django model-based storage for dynamic keywords"""
333
+
334
+ @staticmethod
335
+ def add_keyword(keyword, count=1):
336
+ obj, created = DynamicKeyword.objects.get_or_create(keyword=keyword, defaults={"count": count})
337
+ if not created:
338
+ obj.count += count
339
+ obj.save()
340
+
341
+ @staticmethod
342
+ def get_top_keywords(limit=10):
343
+ return list(DynamicKeyword.objects.order_by("-count").values_list("keyword", flat=True)[:limit])
344
+
345
+ @staticmethod
346
+ def remove_keyword(keyword):
347
+ DynamicKeyword.objects.filter(keyword=keyword).delete()
348
+
349
+ @staticmethod
350
+ def clear_all():
351
+ DynamicKeyword.objects.all().delete()
aiwaf/trainer.py CHANGED
@@ -14,6 +14,7 @@ from django.conf import settings
14
14
  from django.apps import apps
15
15
  from django.db.models import F
16
16
  from .utils import is_exempt_path
17
+ from .storage import get_blacklist_store, get_exemption_store, get_keyword_store
17
18
 
18
19
  # ─────────── Configuration ───────────
19
20
  LOG_PATH = settings.AIWAF_ACCESS_LOG
@@ -28,11 +29,6 @@ _LOG_RX = re.compile(
28
29
  )
29
30
 
30
31
 
31
- BlacklistEntry = apps.get_model("aiwaf", "BlacklistEntry")
32
- DynamicKeyword = apps.get_model("aiwaf", "DynamicKeyword")
33
- IPExemption = apps.get_model("aiwaf", "IPExemption")
34
-
35
-
36
32
  def path_exists_in_django(path: str) -> bool:
37
33
  from django.urls import get_resolver
38
34
  from django.urls.resolvers import URLResolver
@@ -54,13 +50,17 @@ def path_exists_in_django(path: str) -> bool:
54
50
 
55
51
 
56
52
  def remove_exempt_keywords() -> None:
53
+ keyword_store = get_keyword_store()
57
54
  exempt_tokens = set()
55
+
58
56
  for path in getattr(settings, "AIWAF_EXEMPT_PATHS", []):
59
57
  for seg in re.split(r"\W+", path.strip("/").lower()):
60
58
  if len(seg) > 3:
61
59
  exempt_tokens.add(seg)
62
- if exempt_tokens:
63
- DynamicKeyword.objects.filter(keyword__in=exempt_tokens).delete()
60
+
61
+ # Remove exempt tokens from keyword storage
62
+ for token in exempt_tokens:
63
+ keyword_store.remove_keyword(token)
64
64
 
65
65
 
66
66
  def _read_all_logs() -> list[str]:
@@ -98,10 +98,15 @@ def _parse(line: str) -> dict | None:
98
98
 
99
99
  def train() -> None:
100
100
  remove_exempt_keywords()
101
- # Remove any IPs in IPExemption from the blacklist
102
- exempt_ips = set(IPExemption.objects.values_list("ip_address", flat=True))
103
- if exempt_ips:
104
- BlacklistEntry.objects.filter(ip_address__in=exempt_ips).delete()
101
+
102
+ # Remove any IPs in IPExemption from the blacklist using storage system
103
+ exemption_store = get_exemption_store()
104
+ blacklist_store = get_blacklist_store()
105
+
106
+ exempted_ips = [entry['ip_address'] for entry in exemption_store.get_all()]
107
+ for ip in exempted_ips:
108
+ blacklist_store.remove_ip(ip)
109
+
105
110
  raw_lines = _read_all_logs()
106
111
  if not raw_lines:
107
112
  print("No log lines found – check AIWAF_ACCESS_LOG setting.")
@@ -133,10 +138,8 @@ def train() -> None:
133
138
 
134
139
  # Don't block if majority of 404s are on login paths
135
140
  if count > login_404s: # More non-login 404s than login 404s
136
- BlacklistEntry.objects.get_or_create(
137
- ip_address=ip,
138
- defaults={"reason": f"Excessive 404s (≥6 non-login, {count}/{total_404s})"}
139
- )
141
+ blacklist_store = get_blacklist_store()
142
+ blacklist_store.add_ip(ip, f"Excessive 404s (≥6 non-login, {count}/{total_404s})")
140
143
 
141
144
  feature_dicts = []
142
145
  for r in parsed:
@@ -187,10 +190,13 @@ def train() -> None:
187
190
  if anomalous_ips:
188
191
  print(f"⚠️ Detected {len(anomalous_ips)} potentially anomalous IPs during training")
189
192
 
193
+ exemption_store = get_exemption_store()
194
+ blacklist_store = get_blacklist_store()
190
195
  blocked_count = 0
196
+
191
197
  for ip in anomalous_ips:
192
198
  # Skip if IP is exempted
193
- if IPExemption.objects.filter(ip_address=ip).exists():
199
+ if exemption_store.is_exempted(ip):
194
200
  continue
195
201
 
196
202
  # Get this IP's behavior from the data
@@ -213,10 +219,7 @@ def train() -> None:
213
219
  continue
214
220
 
215
221
  # Block if it shows clear signs of malicious behavior
216
- BlacklistEntry.objects.get_or_create(
217
- ip_address=ip,
218
- defaults={"reason": f"AI anomaly + suspicious patterns (kw:{avg_kw_hits:.1f}, 404s:{max_404s}, burst:{avg_burst:.1f})"}
219
- )
222
+ blacklist_store.add_ip(ip, f"AI anomaly + suspicious patterns (kw:{avg_kw_hits:.1f}, 404s:{max_404s}, burst:{avg_burst:.1f})")
220
223
  blocked_count += 1
221
224
  print(f" - {ip}: Blocked for suspicious behavior (kw:{avg_kw_hits:.1f}, 404s:{max_404s}, burst:{avg_burst:.1f})")
222
225
 
@@ -230,8 +233,10 @@ def train() -> None:
230
233
  if len(seg) > 3 and seg not in STATIC_KW:
231
234
  tokens[seg] += 1
232
235
 
233
- for kw, cnt in tokens.most_common(10):
234
- obj, _ = DynamicKeyword.objects.get_or_create(keyword=kw)
235
- DynamicKeyword.objects.filter(pk=obj.pk).update(count=F("count") + cnt)
236
+ keyword_store = get_keyword_store()
237
+ top_tokens = tokens.most_common(10)
238
+
239
+ for kw, cnt in top_tokens:
240
+ keyword_store.add_keyword(kw, cnt)
236
241
 
237
- print(f"DynamicKeyword DB updated with top tokens: {[kw for kw, _ in tokens.most_common(10)]}")
242
+ print(f"DynamicKeyword storage updated with top tokens: {[kw for kw, _ in top_tokens]}")
aiwaf/utils.py CHANGED
@@ -4,7 +4,7 @@ import glob
4
4
  import gzip
5
5
  from datetime import datetime
6
6
  from django.conf import settings
7
- from .models import IPExemption
7
+ from .storage import get_exemption_store
8
8
 
9
9
  _LOG_RX = re.compile(
10
10
  r'(\d+\.\d+\.\d+\.\d+).*\[(.*?)\].*"(GET|POST) (.*?) HTTP/.*?" (\d{3}).*?"(.*?)" "(.*?)"'
@@ -53,7 +53,8 @@ def parse_log_line(line):
53
53
 
54
54
  def is_ip_exempted(ip):
55
55
  """Check if IP is in exemption list"""
56
- return IPExemption.objects.filter(ip_address=ip).exists()
56
+ store = get_exemption_store()
57
+ return store.is_exempted(ip)
57
58
 
58
59
  def is_view_exempt(request):
59
60
  """Check if the current view is marked as AI-WAF exempt"""
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: aiwaf
3
- Version: 0.1.8.6
3
+ Version: 0.1.8.7
4
4
  Summary: AI-powered Web Application Firewall
5
5
  Home-page: https://github.com/aayushgauba/aiwaf
6
6
  Author: Aayush Gauba
@@ -188,6 +188,31 @@ AIWAF_ACCESS_LOG = "/var/log/nginx/access.log"
188
188
 
189
189
  ---
190
190
 
191
+ ### Storage Configuration
192
+
193
+ **Choose storage backend:**
194
+
195
+ ```python
196
+ # Use Django models (default) - requires database tables
197
+ AIWAF_STORAGE_MODE = "models"
198
+
199
+ # OR use CSV files - no database required
200
+ AIWAF_STORAGE_MODE = "csv"
201
+ AIWAF_CSV_DATA_DIR = "aiwaf_data" # Directory for CSV files
202
+ ```
203
+
204
+ **CSV Mode Features:**
205
+ - No database migrations required
206
+ - Files stored in `aiwaf_data/` directory:
207
+ - `blacklist.csv` - Blocked IP addresses
208
+ - `exemptions.csv` - Exempt IP addresses
209
+ - `keywords.csv` - Dynamic keywords
210
+ - `access_samples.csv` - Feature samples for ML training
211
+ - Perfect for lightweight deployments or when you prefer file-based storage
212
+ - Management commands work identically in both modes
213
+
214
+ ---
215
+
191
216
  ### Optional (defaults shown)
192
217
 
193
218
  ```python
@@ -1,22 +1,22 @@
1
1
  aiwaf/__init__.py,sha256=nQFpJ1YpX48snzLjEQCf8zD2YNh8v0b_kPTrXx8uBYc,46
2
2
  aiwaf/apps.py,sha256=nCez-Ptlv2kaEk5HenA8b1pATz1VfhrHP1344gwcY1A,142
3
- aiwaf/blacklist_manager.py,sha256=sM6uTH7zD6MOPGb0kzqV2aFut2vxKgft_UVeRJr7klw,392
3
+ aiwaf/blacklist_manager.py,sha256=92ltIrFfv8WOC4CXwvNVZYfivkRZHGNg3E2QAbHQipQ,550
4
4
  aiwaf/decorators.py,sha256=IUKOdM_gdroffImRZep1g1wT6gNqD10zGwcp28hsJCs,825
5
- aiwaf/middleware.py,sha256=eMad-wvQWALkH2nIhjssU9Y-AqFleP3Gm0lRu3qE0Bw,9679
5
+ aiwaf/middleware.py,sha256=1JPrc0npI_a5bnB-thN0ME1ehfTbWBl1j9wTndZwRdQ,9505
6
6
  aiwaf/models.py,sha256=XaG1pd_oZu3y-fw66u4wblGlWcUY9gvsTNKGD0kQk7Y,1672
7
- aiwaf/storage.py,sha256=bxCILzzvA1-q6nwclRE8WrfoRhe25H4VrsQDf0hl_lY,1903
8
- aiwaf/trainer.py,sha256=R00q_QQ1o2UmdIWMWNh847BGBrnI6j-hfjNalojfnhU,8494
9
- aiwaf/utils.py,sha256=s-rtUrWQFVv-nuGxe2hz5-LLvB6TbZXKj6do46DwrkA,3376
7
+ aiwaf/storage.py,sha256=Z0KWArfLmOHnvUcL5aVx8W_aHMr-qoEW8FVGrM23BvA,11639
8
+ aiwaf/trainer.py,sha256=qzQOtMW0_OA5JWWh6Znbc5BG3gcQ6Cb_NNWMOgLf3VQ,8487
9
+ aiwaf/utils.py,sha256=BJk5vJCYdGPl_4QQiknjhCbkzv5HZCXgFcBJDMJpHok,3390
10
10
  aiwaf/management/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
11
11
  aiwaf/management/commands/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
12
- aiwaf/management/commands/add_ipexemption.py,sha256=LWN21_ydqSjU3_hUnkou4Ciyrk_479zLvcKdWm8hkC0,988
13
- aiwaf/management/commands/aiwaf_reset.py,sha256=dUTYX6Z6_X3Ft3lqF_McXE7OdKADlQFGFWvjdvFVZFI,3245
12
+ aiwaf/management/commands/add_ipexemption.py,sha256=srgdVPDJtF7G9GGIqaZ7L3qTuNheoS_uwlhlRO4W2bc,945
13
+ aiwaf/management/commands/aiwaf_reset.py,sha256=0FIBqpZS8xgFFvAKJ-0zAC_-QNQwRkOHpXb8N-OdFr8,3740
14
14
  aiwaf/management/commands/detect_and_train.py,sha256=-o-LZ7QZ5GeJPCekryox1DGXKMmFEkwwrcDsiM166K0,269
15
15
  aiwaf/resources/model.pkl,sha256=5t6h9BX8yoh2xct85MXOO60jdlWyg1APskUOW0jZE1Y,1288265
16
16
  aiwaf/templatetags/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
17
17
  aiwaf/templatetags/aiwaf_tags.py,sha256=XXfb7Tl4DjU3Sc40GbqdaqOEtKTUKELBEk58u83wBNw,357
18
- aiwaf-0.1.8.6.dist-info/licenses/LICENSE,sha256=Ir8PX4dxgAcdB0wqNPIkw84fzIIRKE75NoUil9RX0QU,1069
19
- aiwaf-0.1.8.6.dist-info/METADATA,sha256=s6gux1GQJsbvphRXyVkVPm63_bVWIXcJcagWtBSlgpE,7955
20
- aiwaf-0.1.8.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
21
- aiwaf-0.1.8.6.dist-info/top_level.txt,sha256=kU6EyjobT6UPCxuWpI_BvcHDG0I2tMgKaPlWzVxe2xI,6
22
- aiwaf-0.1.8.6.dist-info/RECORD,,
18
+ aiwaf-0.1.8.7.dist-info/licenses/LICENSE,sha256=Ir8PX4dxgAcdB0wqNPIkw84fzIIRKE75NoUil9RX0QU,1069
19
+ aiwaf-0.1.8.7.dist-info/METADATA,sha256=qAIF4-sV_Br3BPp6Ivn3fXzY9uWvdtMco1zXBlkKkW0,8664
20
+ aiwaf-0.1.8.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
21
+ aiwaf-0.1.8.7.dist-info/top_level.txt,sha256=kU6EyjobT6UPCxuWpI_BvcHDG0I2tMgKaPlWzVxe2xI,6
22
+ aiwaf-0.1.8.7.dist-info/RECORD,,