#!/usr/bin/env python3
import os
os.nice(20)
import asyncio, subprocess
import arrow
import peewee
import click
from path import Path
import pyufw as ufw
from playhouse.sqlite_ext import SqliteExtDatabase
import yaml
from peewee import fn
from rich import print

base_path = Path('/var/opt/waf')
conf_file = base_path / 'config.yml'

if conf_file.exists():
    conf = yaml.safe_load(conf_file.read_text())
else:
    conf = {
        'vroots': "/srv",
        'logs_glob': "*/logs/*access*.log",
        'whitelist_ips': [
            '127.0.0.1',
        ],
        'date_range': 'last_hour',
        'purge_older_than': 'three_days',
        'db_filename': 'waf.db',
    }
    conf_file.touch()
    yaml.dump(conf, conf_file.open('w'))

now = arrow.utcnow()
older_than = now.shift(days=-3).floor('day')
last_hour = now.shift(hours=-1).floor('hour')
last_thirty_min = now.shift(minutes=-30)

# Configuring
db_path =  base_path / conf['db_filename']
date_range = globals()[conf['date_range']]
logs = Path(conf['vroots']).glob(conf['logs_glob'])
whitelist_ips = conf['whitelist_ips']

# Database

db = SqliteExtDatabase(db_path, pragmas={'journal_mode': 'wal'})

class Attack(peewee.Model):
    ip = peewee.CharField(unique=True)
    host = peewee.CharField(index = True)
    date = peewee.DateTimeField(default=now.datetime)
    count = peewee.BigIntegerField()

    class Meta:
        database = db

Attack.create_table(True)

# Utils

def report_attacks():
    click.echo(
        click.style(
            f"Attacks in database: {Attack.select().count()}",
            fg="cyan"
        )
    )
    hosts = {}
    for a in Attack.select():
        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]))

    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"
            )
        )

# Check rules

checklist = [
    {
        'where': 'url',
        'in': 'xmlrpc',
        'store': 'suspects',
    },
    {
        'where': 'agent',
        'in': 'PHP/6',
        '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)
    if returned_value == 0:
        click.echo(click.style('Nginx reloaded', fg="blue"))

async def get_denied():
    denied = []
    for rule in ufw.get_rules().values():
        if "deny from " in rule:
            denied.append(rule.split(' ')[-1])
    return denied

async def deny(ip):
    ufw.add('deny from ' + ip + ' to any', number=1)

def undeny(ip):
    # Not async to avoid deleting wrong ufw rule.
    for key, rule in ufw.get_rules().items():
        if rule == "deny from " + ip:
            ufw.delete(key)

    # attack = Attack.get_or_none(Attack.ip == ip)
    # if attack:
    #     attack.delete_instance()

async def check(ip, host, date_position):
    date = arrow.get(date_position,'DD/MMM/YYYY:HH:mm:ss')
    if date > date_range:
        find = Attack.get_or_none(Attack.ip == ip)
        if find:
            find.count = find.count + 1
            find.save()
        else:
            data = {'ip': ip, 'date':date.datetime, 'host': host, 'count':1}
            Attack.create(**data)

async def scan(log):
    suspects = []
    suspects_404 = {}
    for line in log.lines():
        splitted = line.split()
        # print(splitted)
        # print(splitted[-2].strip())
        ip = splitted[0].strip()
        method = splitted[5].strip()[1:]
        url = splitted[6].strip()
        agent = splitted[-2].strip()
        status = splitted[8].strip()
        date_position = splitted[3][1:]
        host = log.splitall()[2]
        if ip not in whitelist_ips:
            for rule in checklist:
                where = url
                if rule['where'] == 'url':
                    where = url
                elif rule['where'] == 'agent':
                    where = agent
                store = suspects
                if rule['store'] == 'suspects':
                    store = suspects
                if 'in' in rule and rule['in'] in where:
                    # if rule['where'] == 'agent':
                    #     print(where)
                    store.append(check(ip, host, date_position))
                    break
                elif 'startswith' in rule and url.startswith(rule['startswith']) and 'notin' in rule and rule['notin'] not in url:
                    store.append(check(ip, host, date_position))
                    break
                elif 'startswith' in rule and url.startswith(rule['startswith']):
                    store.append(check(ip, host, date_position))
                    break
            # elif 'wp-admin' in url and status not in ['200','302','499']:
            #     store.append(check(ip, host, date_position))

    # def is_suspicious_login(item):
    #     return len(item[1]) > 18
    # filtered = dict(filter(is_suspicious_login, suspects_login.items()))

    # for ip,suspect in filtered.items():
    #     suspects.append(check(ip, suspect[-1][1], suspect[-1][2]))

    await asyncio.gather(*suspects)



async def block():
    denied = await get_denied()
    found = Attack.select().where(
        (Attack.ip.not_in(denied)) &
        (Attack.count > 15)
    )
    if found.count() > 0:
        click.echo(click.style('New IPs to block: {}'.format(found.count()), fg="yellow"))
        to_deny = []
        for attack in found:
            to_deny.append( deny(attack.ip) )
        await asyncio.gather(*to_deny)
        await nginx_reload()
    else:
        click.echo(click.style('No IPs to block', fg="blue"))



async def start():
    report()
    scans = []
    for log in logs:
        scans.append(scan(log))
    await asyncio.gather(*scans)
    await block()
    report_attacks()
    
async def info(ip, unblock=False):
    denied = await get_denied()
    find = Attack.get_or_none(Attack.ip == ip)
    if find:
        click.echo(click.style(f'IP {find.ip} found with {find.count} attacks at {find.host}', fg="yellow", bold=True))
        if unblock is True:
            find.delete_instance()
            click.echo(click.style(f'IP {ip} removed from attacks', fg="green", bold=True))
    else:
        click.echo(click.style(f'IP {ip} no attacks found', fg="blue"))

    if ip in denied:
        if unblock is False:
            click.echo(click.style(f'IP {ip} is BLOCKED', fg="yellow", bold=True))
        else:
            undeny(ip)
            click.echo(click.style(f'IP {ip} UNBLOCKED', fg="green", bold=True))


def purge():
    found = Attack.select(Attack.ip).where(
        (Attack.date < older_than.datetime ) )
    for attack in found:
        undeny(attack.ip)
        attack.delete_instance()

def empty():
    found = Attack.select()
    for attack in found:
        undeny(attack.ip)
        attack.delete_instance()

# CLI App
@click.group()
def cli():
    pass

@cli.command('info')
@click.argument('ip')
def waf_info(ip):
    click.echo('Find blocked IP')
    asyncio.run(info(ip))


@cli.command('unblock')
@click.argument('ip')
def waf_unblock(ip):
    click.echo('Unblock blocked IP')
    asyncio.run(info(ip,unblock=True))


@cli.command('purge')
def waf_purge():
    click.echo('Clean old blocked IPs')
    purge()


@cli.command('empty')
def waf_empty():
    click.echo('Clean all blocked IPs')
    empty()


@cli.command('report')
def waf_report():
    click.echo(click.style('Report WAF attacks', fg="blue", bold=True))
    report()


@cli.command('scan')
def waf_scan():
    before = arrow.utcnow()
    click.echo(click.style('Scan', fg="blue", blink=True, bold=True))
    asyncio.run(start())
    after = arrow.utcnow()
    click.echo(click.style(f'Finished in {after - before}', fg="blue"))


if __name__ == '__main__':
    cli()