#!/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()