diff --git a/beam.py b/beam.py index 71f4355..1702269 100644 --- a/beam.py +++ b/beam.py @@ -29,6 +29,9 @@ def __init__(self, debug="INFO", **kwargs): self.message_id = 0 self.csrf_token = None + self.quiet = None + self.bot_id = None + self.channel_id = None self.http_session = Session() @@ -136,10 +139,14 @@ def get_chat_users(self, id): def connect(self, channel_id, bot_id, quiet=False): """Connect to a Beam chat through a websocket.""" + self.channel_id = channel_id + self.bot_id = bot_id + self.quiet = quiet + self.connection_information = { - "channel_id": channel_id, - "bot_id": bot_id, - "quiet": quiet + "channel_id": self.channel_id, + "bot_id": self.bot_id, + "quiet": self.quiet } chat = self.get_chat(channel_id) @@ -165,8 +172,6 @@ def connect(self, channel_id, bot_id, quiet=False): def authenticate(self, *args): """Authenticate session to a Beam chat through a websocket.""" - backoff = 0 - future = args[-1] if future.exception() is None: self.websocket = future.result() @@ -181,12 +186,11 @@ def authenticate(self, *args): self.read_chat(self.handle) else: self.logger.error("There was an issue connecting.") - self.logger.error("Trying again in {} seconds.".format(backoff)) + self.logger.error("Trying again in 10 seconds.") - time.sleep(min(2**backoff, 60)) - backoff += 1 + time.sleep(10) - self.authenticate(*args) + self.connect(self.channel_id, self.bot_id, self.quiet) def send_message(self, *args, method="msg"): """Send a message to a Beam chat through a websocket.""" diff --git a/cactus.py b/cactus.py index 974b914..1797949 100644 --- a/cactus.py +++ b/cactus.py @@ -54,7 +54,6 @@ def __init__(self, **kwargs): self.debug = kwargs.get("debug", False) self.config_file = kwargs.get("config_file", "data/config.json") - self.stats_file = kwargs.get("stats_file", "data/stats.json") self.database = kwargs.get("database", "data/data.db") self.quiet = kwargs.get("quiet", False) @@ -90,23 +89,6 @@ def load_config(self, filename): raise FileNotFoundError("Configuration file not found.") exit() - def load_stats(self, filename): - """Load statistics file.""" - - self.logger.warning("Statistics are not yet implemented.") - return dict() - - if exists(filename): - self.stats_file = filename - self.logger.info("Statistics file found. Loading...") - with open(filename) as stats: - self.stats = load(stats) - return self.stats - else: - self.logger.warn("Statistics file not found. Creating...") - copyfile("data/stats-template.json", "data/stats.json") - self.logger.info("Statistics file created.") - def update_config(self, keys, value): """Update configuration file value.""" @@ -119,33 +101,18 @@ def update_config(self, keys, value): self.config = config_data return self.config - def update_stats(self, keys, value): - """Update statistics file value.""" - - self.logger.warning("Statistics are not yet implemented.") - return - - with open(self.stats_file, 'r') as stats: - stats_data = load(stats) - reduce(lambda d, k: d[k], keys.split('.')[:-1], stats_data)[ - keys.split('.')[-1]] = value - with open(self.stats_file, 'w+') as stats: - dump(stats_data, stats, indent=2, sort_keys=True) - self.stats = stats_data - return self.stats - def run(self, *args, **kwargs): """Run bot.""" self.logger.info(cactus_art) self._init_database(self.database) self.load_config(filename=self.config_file) - self.load_stats(filename=self.stats_file) self.started = True while self.config.get("autorestart") or not self.started: try: self.bot_data = self.login(**self.config["auth"]) + self.bot_name = self.config["auth"]["username"] self.logger.info("Authenticated as: {}.".format( self.bot_data["username"])) diff --git a/data/stats-template.json b/data/stats-template.json deleted file mode 100644 index 1e3ad54..0000000 --- a/data/stats-template.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "status": "Not yet implemented." -} diff --git a/messages.py b/messages.py index 8092735..dbe3278 100644 --- a/messages.py +++ b/messages.py @@ -5,8 +5,8 @@ from beam import Beam from models import (Command, User, session, CommandCommand, QuoteCommand, CubeCommand, SocialCommand, UptimeCommand, PointsCommand, - TemmieCommand, FriendCommand, SpamProtCommand, ProCommand, - SubCommand, RepeatCommand) + TemmieCommand, FriendCommand, SpamProtCommand, + RepeatCommand) class MessageHandler(Beam): @@ -38,8 +38,6 @@ def _init_commands(self): "friend": FriendCommand(self.get_channel), "points": PointsCommand(self.config["points"]["name"]), "spamprot": SpamProtCommand(self.update_config), - "pro": ProCommand(), - "sub": SubCommand(), "cube": CubeCommand(), "temmie": TemmieCommand() } @@ -89,7 +87,8 @@ def message_handler(self, data): session.commit() mod_roles = ("Owner", "Staff", "Founder", "Global Mod", "Mod") - if not (data["user_roles"][0] in mod_roles or user.friend): + if not (data["user_roles"][0] in mod_roles or user.friend or + self.bot_name.lower() == data["user_name"].lower()): if (len(parsed) > self.config["spam_protection"].get( "maximum_message_length", 256)): self.remove_message(data["id"]) diff --git a/models.py b/models.py index 7070696..06a279e 100644 --- a/models.py +++ b/models.py @@ -491,20 +491,22 @@ def __call__(self, args, data): if len(args) < 2: return "Not enough arguments." elif len(args) == 2: - match = re.match(r'@?([A-Za-z0-9]{,32})', args[1]) - if match is None: + user = re.match(r'@?([\w_-]*[a-z][\w_-]*', args[1]) + if user is None: return "Invalid username '{}'.".format(args[1]) - id = self.get_channel(args[1]) - if id.get("statusCode") == 404: + channel_id = self.get_channel(user.group()) + + if channel_id.get("statusCode") == 404: return "User has not entered this channel." - query = session.query(User).filter_by(id=id["user"]["id"]).first() + query = session.query(User).filter_by( + id=channel_id["user"]["id"]).first() if query: query.friend = not query.friend session.commit() return "{}ed @{} as a friend.".format( - ["Remov", "Add"][query.friend], args[1]) + ["Remov", "Add"][query.friend], user.group()) else: return "User has not entered this channel." elif len(args) > 2: @@ -557,17 +559,3 @@ def __call__(self, args, data=None): return "Invalid true/false: '{}'.".format(args[2]) return "Invalid argument: '{}'.".format(args[1]) return "Not enough arguments." - - -class ProCommand(Command): - - @role_specific("Pro", reply="pro") - def __call__(self, args=None, data=None): - return "I'm such a Pro! B)" - - -class SubCommand(Command): - - @role_specific("Subscriber", reply="sub") - def __call__(self, args=None, data=None): - return "I'm a subscriber! :salute"