コード例 #1
0
def test_bucket_regen():
    bucket = TokenBucket(10, 10)
    # success
    assert bucket.consume(10) is True
    # sleep
    time.sleep(1)
    # bucket should be full again and this should succeed
    assert bucket.tokens == 10
    assert bucket.consume(10) is True
コード例 #2
0
def test_bucket_consume():
    bucket = TokenBucket(10, 5)
    # larger then capacity
    assert bucket.consume(15) is False
    # success
    assert bucket.consume(10) is True
    # check if bucket has no tokens
    assert bucket._tokens == 0
    # bucket is empty from above, should fail
    assert bucket.consume(10) is False
コード例 #3
0
def test_bucket_advanced():
    bucket = TokenBucket(10, 1)
    # tokens start at 10
    assert bucket._tokens == 10
    # empty tokens
    assert bucket.empty() is True
    # check tokens is 0
    assert bucket._tokens == 0
    # refill tokens
    assert bucket.refill() is True
    # check tokens is 10
    assert bucket._tokens == 10
コード例 #4
0
def rate_limit(bot: CloudBot, event: Event, _hook: Hook) -> Optional[Event]:
    """
    Handle rate limiting certain hooks
    """
    conn = event.conn
    # check command spam tokens
    if _hook.type in ("command", "regex"):
        uid = "!".join([conn.name, event.chan, event.nick]).lower()

        config = conn.config.get("ratelimit", {})
        tokens = config.get("tokens", 17.5)
        restore_rate = config.get("restore_rate", 2.5)
        message_cost = config.get("message_cost", 5)
        strict = config.get("strict", True)

        try:
            bucket = buckets[uid]
        except KeyError:
            buckets[uid] = bucket = TokenBucket(tokens, restore_rate)

        if not bucket.consume(message_cost):
            logger.info(
                "[%s|sieve] Refused command from %s. Entity had %s tokens, needed %s.",
                conn.name,
                uid,
                bucket.tokens,
                message_cost,
            )
            if strict:
                # bad person loses all tokens
                bucket.empty()

            return None

    return event
コード例 #5
0
def test_task_clear() -> None:
    core_sieve.buckets["a"] = bucket = TokenBucket(10, 2)
    bucket.timestamp = 0
    assert len(core_sieve.buckets) == 1
    core_sieve.task_clear()
    assert len(core_sieve.buckets) == 0
コード例 #6
0
ファイル: core_sieve.py プロジェクト: oyvindio/CloudBot
def sieve_suite(bot, event, _hook):
    """
    this function stands between your users and the commands they want to use. it decides if they can or not
    :type bot: cloudbot.bot.CloudBot
    :type event: cloudbot.event.Event
    :type _hook: cloudbot.plugin.Hook
    """
    conn = event.conn
    # check ignore bots
    if event.irc_command == 'PRIVMSG' and event.nick.endswith('bot') and _hook.ignore_bots:
        return None

    # check acls
    acl = conn.config.get('acls', {}).get(_hook.function_name)
    if acl:
        if 'deny-except' in acl:
            allowed_channels = list(map(str.lower, acl['deny-except']))
            if event.chan.lower() not in allowed_channels:
                return None
        if 'allow-except' in acl:
            denied_channels = list(map(str.lower, acl['allow-except']))
            if event.chan.lower() in denied_channels:
                return None

    # check disabled_commands
    if _hook.type == "command":
        disabled_commands = conn.config.get('disabled_commands', [])
        if event.triggered_command in disabled_commands:
            return None

    # check permissions
    allowed_permissions = _hook.permissions
    if allowed_permissions:
        allowed = False
        for perm in allowed_permissions:
            if event.has_permission(perm):
                allowed = True
                break

        if not allowed:
            event.notice("Sorry, you are not allowed to use this command.")
            return None

    # check command spam tokens
    if _hook.type == "command":
        # right now ratelimiting is per-channel, but this can be changed
        uid = (event.chan, event.nick.lower())

        if uid not in buckets:
            bucket = TokenBucket(TOKENS, RESTORE_RATE)
            bucket.consume(MESSAGE_COST)
            buckets[uid] = bucket
            return event

        bucket = buckets[uid]
        if bucket.consume(MESSAGE_COST):
            pass
        else:
            bot.logger.info("[{}|sieve] Refused command from {}. "
                            "Entity had {} tokens, needed {}.".format(conn.readable_name, uid, bucket.tokens,
                                                                      MESSAGE_COST))
            if STRICT:
                # bad person loses all tokens
                bucket.empty()
            return None

    return event
コード例 #7
0
async def sieve_suite(bot, event, _hook):
    conn = event.conn

    # check acls
    acl = conn.config.get('acls', {}).get(_hook.function_name)
    if acl:
        if 'deny-except' in acl:
            allowed_channels = list(map(str.lower, acl['deny-except']))
            if event.chan.lower() not in allowed_channels:
                return None
        if 'allow-except' in acl:
            denied_channels = list(map(str.lower, acl['allow-except']))
            if event.chan.lower() in denied_channels:
                return None

    # check disabled_commands
    if _hook.type == "command":
        disabled_commands = conn.config.get('disabled_commands', [])
        if event.triggered_command in disabled_commands:
            return None

    # check permissions
    allowed_permissions = _hook.permissions
    if allowed_permissions:
        allowed = False
        for perm in allowed_permissions:
            if await event.check_permission(perm):
                allowed = True
                break

        if not allowed:
            event.notice("Sorry, you are not allowed to use this command.")
            return None

    # check command spam tokens
    if _hook.type == "command":
        uid = "!".join([conn.name, event.chan, event.nick]).lower()

        tokens = conn.config.get('ratelimit', {}).get('tokens', 17.5)
        restore_rate = conn.config.get('ratelimit', {}).get('restore_rate', 2.5)
        message_cost = conn.config.get('ratelimit', {}).get('message_cost', 5)
        strict = conn.config.get('ratelimit', {}).get('strict', True)

        if uid not in buckets:
            bucket = TokenBucket(tokens, restore_rate)
            bucket.consume(message_cost)
            buckets[uid] = bucket
            return event

        bucket = buckets[uid]
        if bucket.consume(message_cost):
            pass
        else:
            logger.info(
                "[%s|sieve] Refused command from %s. "
                "Entity had %s tokens, needed %s.",
                conn.name, uid, bucket.tokens, message_cost
            )
            if strict:
                # bad person loses all tokens
                bucket.empty()
            return None

    return event
コード例 #8
0
ファイル: core_sieve.py プロジェクト: ps00k/WeedBot
def sieve_suite(bot, event, _hook):
    """
    this function stands between your users and the commands they want to use. it decides if they can or not
    :type bot: cloudbot.bot.CloudBot
    :type event: cloudbot.event.Event
    :type _hook: cloudbot.plugin.Hook
    """
    conn = event.conn
    # check ignore bots
    if event.irc_command == 'PRIVMSG' and event.nick.endswith(
            'bot') and _hook.ignore_bots:
        return None

    # check acls
    acl = conn.config.get('acls', {}).get(_hook.function_name)
    if acl:
        if 'deny-except' in acl:
            allowed_channels = list(map(str.lower, acl['deny-except']))
            if event.chan.lower() not in allowed_channels:
                return None
        if 'allow-except' in acl:
            denied_channels = list(map(str.lower, acl['allow-except']))
            if event.chan.lower() in denied_channels:
                return None

    # check disabled_commands
    if _hook.type == "command":
        disabled_commands = conn.config.get('disabled_commands', [])
        if event.triggered_command in disabled_commands:
            return None

    # check permissions
    allowed_permissions = _hook.permissions
    if allowed_permissions:
        allowed = False
        for perm in allowed_permissions:
            if event.has_permission(perm):
                allowed = True
                break

        if not allowed:
            event.notice("Sorry, you are not allowed to use this command.")
            return None

    # check command spam tokens
    if _hook.type == "command":
        # right now ratelimiting is per-channel, but this can be changed
        uid = (event.chan, event.nick.lower())

        tokens = conn.config.get('ratelimit', {}).get('tokens', 17.5)
        restore_rate = conn.config.get('ratelimit',
                                       {}).get('restore_rate', 2.5)
        message_cost = conn.config.get('ratelimit', {}).get('message_cost', 5)
        strict = conn.config.get('ratelimit', {}).get('strict', True)

        if uid not in buckets:
            bucket = TokenBucket(tokens, restore_rate)
            bucket.consume(message_cost)
            buckets[uid] = bucket
            return event

        bucket = buckets[uid]
        if bucket.consume(message_cost):
            pass
        else:
            bot.logger.info("[{}|sieve] Refused command from {}. "
                            "Entity had {} tokens, needed {}.".format(
                                conn.name, uid, bucket.tokens, message_cost))
            if strict:
                # bad person loses all tokens
                bucket.empty()
            return None

    return event
コード例 #9
0
ファイル: core_sieve.py プロジェクト: Abrackadabra/CloudBot
def sieve_suite(bot, event, _hook):
    global buckets

    conn = event.conn

    # check acls
    acl = conn.config.get('acls', {}).get(_hook.function_name)
    if acl:
        if 'deny-except' in acl:
            allowed_channels = list(map(str.lower, acl['deny-except']))
            if event.chan.lower() not in allowed_channels:
                return None
        if 'allow-except' in acl:
            denied_channels = list(map(str.lower, acl['allow-except']))
            if event.chan.lower() in denied_channels:
                return None

    # check disabled_commands
    if _hook.type == "command":
        disabled_commands = conn.config.get('disabled_commands', [])
        if event.triggered_command in disabled_commands:
            return None

    # check permissions
    allowed_permissions = _hook.permissions
    if allowed_permissions:
        allowed = False
        for perm in allowed_permissions:
            if event.has_permission(perm):
                allowed = True
                break

        if not allowed:
            event.notice("Sorry, you are not allowed to use this command.")
            return None

    # check command spam tokens
    if _hook.type == "command":
        uid = "!".join([conn.name, event.chan, event.nick]).lower()

        tokens = conn.config.get('ratelimit', {}).get('tokens', 17.5)
        restore_rate = conn.config.get('ratelimit', {}).get('restore_rate', 2.5)
        message_cost = conn.config.get('ratelimit', {}).get('message_cost', 5)
        strict = conn.config.get('ratelimit', {}).get('strict', True)

        if uid not in buckets:
            bucket = TokenBucket(tokens, restore_rate)
            bucket.consume(message_cost)
            buckets[uid] = bucket
            return event

        bucket = buckets[uid]
        if bucket.consume(message_cost):
            pass
        else:
            bot.logger.info("[{}|sieve] Refused command from {}. "
                            "Entity had {} tokens, needed {}.".format(conn.name, uid, bucket.tokens,
                                                                      message_cost))
            if strict:
                # bad person loses all tokens
                bucket.empty()
            return None

    return event
コード例 #10
0
ファイル: core_sieve.py プロジェクト: skius/CloudBot
def sieve_suite(bot, event, _hook):
    global buckets

    conn = event.conn

    # check acls
    acl = conn.config.get('acls', {}).get(_hook.function_name)
    if acl:
        if 'deny-except' in acl:
            allowed_channels = list(map(str.lower, acl['deny-except']))
            if event.chan.lower() not in allowed_channels:
                return None
        if 'allow-except' in acl:
            denied_channels = list(map(str.lower, acl['allow-except']))
            if event.chan.lower() in denied_channels:
                return None

    # check disabled_commands
    if _hook.type == "command":
        disabled_commands = conn.config.get('disabled_commands', [])
        if event.triggered_command in disabled_commands:
            return None

    # check permissions
    allowed_permissions = _hook.permissions
    if allowed_permissions:
        allowed = False
        for perm in allowed_permissions:
            if (yield from event.check_permission(perm)):
                allowed = True
                break

        if not allowed:
            event.reply("I'm sorry, Dave. I'm afraid I can't do that.")
            return None

    # check command spam tokens
    if _hook.type == "command":
        uid = "!".join([conn.name, event.chan, event.nick]).lower()

        tokens = conn.config.get('ratelimit', {}).get('tokens', 17.5)
        restore_rate = conn.config.get('ratelimit',
                                       {}).get('restore_rate', 2.5)
        message_cost = conn.config.get('ratelimit', {}).get('message_cost', 5)
        strict = conn.config.get('ratelimit', {}).get('strict', True)

        if uid not in buckets:
            bucket = TokenBucket(tokens, restore_rate)
            bucket.consume(message_cost)
            buckets[uid] = bucket
            return event

        bucket = buckets[uid]
        if bucket.consume(message_cost):
            pass
        else:
            bot.logger.info("[{}|sieve] Refused command from {}. "
                            "Entity had {} tokens, needed {}.".format(
                                conn.name, uid, bucket.tokens, message_cost))
            if strict:
                # bad person loses all tokens
                bucket.empty()
            return None

    return event
コード例 #11
0
ファイル: core_sieve.py プロジェクト: Vault108-zz/CB
def sieve_suite(bot, event, _hook):
    global buckets

    conn = event.conn

    # check acls
    acl = conn.config.get('acls', {}).get(_hook.function_name)
    if acl:
        if 'deny-except' in acl:
            allowed_channels = list(map(str.lower, acl['deny-except']))
            if event.chan.lower() not in allowed_channels:
                return None
        if 'allow-except' in acl:
            denied_channels = list(map(str.lower, acl['allow-except']))
            if event.chan.lower() in denied_channels:
                return None

    # check disabled_commands
    if _hook.type == "command":
        disabled_commands = conn.config.get('disabled_commands', [])
        if event.triggered_command in disabled_commands:
            return None

    # check permissions
    allowed_permissions = _hook.permissions
    if allowed_permissions:
        allowed = False
        for perm in allowed_permissions:
            if event.has_permission(perm):
                allowed = True
                break

            for perm_hook in bot.plugin_manager.perm_hooks[perm]:
                try:
                    res = yield from async_util.run_func(
                        event.loop, perm_hook.function, bot, event, _hook)
                except Exception:
                    logger.exception("Error in hook {}".format(
                        perm_hook.description))
                else:
                    if res:
                        allowed = True
                        break

            if allowed:
                break

        if not allowed:
            event.notice("Sorry, you are not allowed to use this command.")
            return None

    # check command spam tokens
    if _hook.type == "command":
        uid = "!".join([conn.name, event.chan, event.nick]).lower()

        tokens = conn.config.get('ratelimit', {}).get('tokens', 17.5)
        restore_rate = conn.config.get('ratelimit',
                                       {}).get('restore_rate', 2.5)
        message_cost = conn.config.get('ratelimit', {}).get('message_cost', 5)
        strict = conn.config.get('ratelimit', {}).get('strict', True)

        if uid not in buckets:
            bucket = TokenBucket(tokens, restore_rate)
            bucket.consume(message_cost)
            buckets[uid] = bucket
            return event

        bucket = buckets[uid]
        if bucket.consume(message_cost):
            pass
        else:
            bot.logger.info("[{}|sieve] Refused command from {}. "
                            "Entity had {} tokens, needed {}.".format(
                                conn.name, uid, bucket.tokens, message_cost))
            if strict:
                # bad person loses all tokens
                bucket.empty()
            return None

    return event