rangeban.py

#!/usr/bin/env python3

from datetime import datetime
from ipaddress import ip_address
import sys

import click

from nyaa import create_app, models
from nyaa.extensions import db


def is_cidr_valid(c):
    """Checks whether a CIDR range string is valid."""
    try:
        subnet, mask = c.split("/")
    except ValueError:
        return False
    if int(mask) < 1 or int(mask) > 32:
        return False
    try:
        ip = ip_address(subnet)
    except ValueError:
        return False
    return True


def check_str(b):
    """Returns a checkmark or cross depending on the condition."""
    return "\u2713" if b else "\u2717"


@click.group()
def rangeban():
    global app
    app = create_app("config")


@rangeban.command()
@click.option(
    "--temp/--no-temp",
    help="Mark this entry as one that may be " "cleaned out occasionally.",
    default=False,
)
@click.argument("cidrrange")
def ban(temp, cidrrange):
    if not is_cidr_valid(cidrrange):
        click.secho(
            "{} is not of the format xxx.xxx.xxx.xxx/xx.".format(cidrrange),
            err=True,
            fg="red",
        )
        sys.exit(1)
    with app.app_context():
        ban = models.RangeBan(
            cidr_string=cidrrange, temp=datetime.utcnow() if temp else None
        )
        db.session.add(ban)
        db.session.commit()
        click.echo("Added {} for {}.".format("temp ban" if temp else "ban", cidrrange))


@rangeban.command()
@click.argument("cidrrange")
def unban(cidrrange):
    if not is_cidr_valid(cidrrange):
        click.secho(
            "{} is not of the format xxx.xxx.xxx.xxx/xx.".format(cidrrange),
            err=True,
            fg="red",
        )
        sys.exit(1)
    with app.app_context():
        # Dunno why this wants _cidr_string and not cidr_string, probably
        # due to this all being a janky piece of shit.
        bans = models.RangeBan.query.filter(
            models.RangeBan._cidr_string == cidrrange
        ).all()
        if len(bans) == 0:
            click.echo("Ban not found.")
        for b in bans:
            click.echo("Unbanned {}".format(b.cidr_string))
            db.session.delete(b)
        db.session.commit()


@rangeban.command()
def list():
    with app.app_context():
        bans = models.RangeBan.query.all()
        if len(bans) == 0:
            click.echo("No bans.")
        else:
            click.secho("ID     CIDR Range         Enabled Temp", bold=True)
            for b in bans:
                click.echo(
                    "{0: <6} {1: <18} {2: <7} {3: <4}".format(
                        b.id,
                        b.cidr_string,
                        check_str(b.enabled),
                        check_str(b.temp is not None),
                    )
                )


@rangeban.command()
@click.argument("banid", type=int)
@click.argument("status")
def enabled(banid, status):
    yeses = ["true", "1", "yes", "\u2713"]
    noses = ["false", "0", "no", "\u2717"]
    if status.lower() in yeses:
        set_to = True
    elif status.lower() in noses:
        set_to = False
    else:
        click.secho(
            "Please choose one of {} or {}.".format(yeses, noses), err=True, fg="red"
        )
        sys.exit(1)
    with app.app_context():
        ban = models.RangeBan.query.get(banid)
        if not ban:
            click.secho("No ban with id {} found.".format(banid), err=True, fg="red")
            sys.exit(1)
        ban.enabled = set_to
        db.session.add(ban)
        db.session.commit()
        click.echo(
            "{} ban {} on {}.".format(
                "Enabled" if set_to else "Disabled", banid, ban._cidr_string
            )
        )


if __name__ == "__main__":
    rangeban()