restyle

waf3.py

@@ -45,6 +45,7 @@ logs = vroots.glob('*/logs/*access*.log')
 whitelist_ips = conf['whitelist_ips']
 
 # Database
+
 db = SqliteExtDatabase(db_path, pragmas={'journal_mode': 'wal'})
 
 class Attack(peewee.Model):
@@ -58,10 +59,62 @@ class Attack(peewee.Model):
 
 Attack.create_table(True)
 
-# CLI App
-@click.group()
-def cli():
-    pass
+# Utils
+def report():
+    click.echo(
+        click.style(
+            f"Config file at: {conf_file}",
+            fg="cyan"
+        )
+    )
+    click.echo(
+        click.style(
+            f"Hosting logs: {len(logs)}",
+            fg="cyan"
+        )
+    )
+    report_attacks()
+    for ip in whitelist_ips:
+        click.echo(
+            click.style(
+                f"Whitelisted: {ip}",
+                fg="green"
+            )
+        )
+
+# Check rules
+
+checklist = [
+    {
+        'where': 'url',
+        'in': 'xmlrpc',
+        'store': 'suspects',
+    },
+    {
+        'where': 'url',
+        'in': "shell",
+        'store': 'suspects',
+    },
+    {
+        'where': 'url',
+        'in': "\\x00",
+        'store': 'suspects',
+    },
+    {
+        'method': 'post',
+        'where': 'url',
+        'in': 'wp-login',
+        'store': 'suspects',
+    },
+    {
+        'where': 'url',
+        'startswith': '/.',
+        'notin': '.well_known',
+        'store': 'suspects',
+    },
+]
+
+# App
 
 async def nginx_reload():
     returned_value = subprocess.call('/usr/bin/systemctl reload nginx', shell=True)
@@ -99,40 +152,8 @@ async def check(ip, host, date_position):
             data = {'ip': ip, 'date':date.datetime, 'host': host, 'count':1}
             Attack.create(**data)
 
-checklist = [
-    {
-        'where': 'url',
-        'in': 'xmlrpc',
-        'store': 'suspects',
-    },
-    {
-        'where': 'url',
-        'in': "shell",
-        'store': 'suspects',
-    },
-    {
-        'where': 'url',
-        'in': "\\x00",
-        'store': 'suspects',
-    },
-    {
-        'method': 'post',
-        'where': 'url',
-        'in': 'wp-login',
-        'store': 'suspects',
-    },
-    {
-        'where': 'url',
-        'startswith': '/.',
-        'notin': '.well_known',
-        'store': 'suspects',
-    },
-]
-
-
 async def scan(log):
     suspects = []
-    # suspects_login = {}
     suspects_404 = {}
     for line in log.lines():
         splitted = line.split()
@@ -190,55 +211,6 @@ async def block():
         click.echo(click.style('No IPs to block', fg="blue"))
 
 
-def report_attacks():
-    click.echo(
-        click.style(
-            f"Attacks in database: {Attack.select().count()}",
-            fg="cyan"
-        )
-    )
-    hosts = {}
-    for a in Attack.select():
-        # print(a.host)
-        if a.host in hosts:
-            hosts[a.host] = hosts[a.host] + 1
-        else:
-            hosts[a.host] = 1
-
-    sorted_hosts = dict(sorted(hosts.items(), key=lambda x:x[1]))
-
-    # print(sorted_footballers_by_goals)
-    # query = (Attack
-    #      .select( Attack.id, Attack.host, fn.SUM(Attack.host).alias('sum_host') )
-    #      .group_by(Attack.host)
-    #      .order_by( fn.SUM(Attack.host).alias('sum_host') ))
-
-    for h, v in sorted_hosts.items():
-        print(h, v)
-
-
-def report():
-    click.echo(
-        click.style(
-            f"Config file at: {conf_file}",
-            fg="cyan"
-        )
-    )
-    click.echo(
-        click.style(
-            f"Hosting logs: {len(logs)}",
-            fg="cyan"
-        )
-    )
-    report_attacks()
-    for ip in whitelist_ips:
-        click.echo(
-            click.style(
-                f"Whitelisted: {ip}",
-                fg="green"
-            )
-        )
-
 
 async def start():
     report()
@@ -248,9 +220,7 @@ async def start():
     await asyncio.gather(*scans)
     await block()
     report_attacks()
-    #print("Updated number of attacks in database: {}".format( Attack.select().count() ) )
-
-
+    
 async def info(ip, unblock=False):
     denied = await get_denied()
     find = Attack.get_or_none(Attack.ip == ip)
@@ -278,14 +248,16 @@ def purge():
         undeny(attack.ip)
         attack.delete_instance()
 
-
 def empty():
     found = Attack.select()
     for attack in found:
-        # print(attack.ip)
         undeny(attack.ip)
         attack.delete_instance()
 
+# CLI App
+@click.group()
+def cli():
+    pass
 
 @cli.command('info')
 @click.argument('ip')
@@ -323,7 +295,6 @@ def waf_report():
 def waf_scan():
     before = arrow.utcnow()
     click.echo(click.style('Scan', fg="blue", blink=True, bold=True))
-    # click.echo('Scan')
     asyncio.run(start())
     after = arrow.utcnow()
     click.echo(click.style(f'Finished in {after - before}', fg="blue"))