nxy/bot/plugins/rape.py

# -*- coding: utf-8 -*-
import random

import irc3
from docopt import Dict
from irc3.plugins.command import command
from irc3.utils import IrcString
from psycopg2 import Error

from . import DatabasePlugin


@irc3.plugin
class Rape(DatabasePlugin):
    requires = ['irc3.plugins.command',
                'bot.plugins.storage']

    @command
    def owe(self, mask: IrcString, target: IrcString, args: Dict):
        """Shows how much a nick owes.

        %%owe [<nick>]
        """
        nick = args.get('<nick>', mask.nick)

        # Fetch result from database
        self.cur.execute('''
        select
          fines
        from
          users
        where
          lower(nick) = lower(%s)
        ''', [nick])
        owes = self.cur.fetchone()

        # Colorize owe amount and return string
        if owes:
            fines = '4${fines}'.format(fines=owes['fines'])
        else:
            fines = '3$0'

        # Return total owes
        return '{nick} owes: \x03{fines}\x03'.format(nick=nick, fines=fines)

    @command
    def rape(self, mask: IrcString, target: IrcString, args: Dict):
        """Rapes a nick and eventually charge for it.

        %%rape <nick>
        """
        nick = args.get('<nick>', mask.nick)
        rand = random.randint(0, 3)

        if rand not in (0, 1):
            self.bot.action(target, 'rapes {nick}'.format(nick=nick))
        else:
            fine = random.randint(1, 500)

            try:
                # Insert or add fine to database and return total owe
                self.cur.execute('''
                insert into
                  users (nick, fines)
                values
                  (lower(%s), %s)
                on conflict (nick) do update set
                  fines = users.fines + excluded.fines
                returning
                  fines
                ''', [mask.nick, fine])
                self.con.commit()

                # Get reason based on rand value
                reason = ('raping', 'being too lewd and getting raped')[rand]

                # Print fine and total owe
                self.bot.action(target,
                                'fines {nick} \x02${fine}\x02 for {reason}. '
                                'You owe: \x0304${total}\x03'
                                .format(nick=mask.nick,
                                        fine=fine,
                                        reason=reason,
                                        total=self.cur.fetchone()['fines']))
            except Error:
                # Rollback transaction on error
                self.con.rollback()