From 1d3b541e8f8e56a39f32e9aaea819743ae7a24be Mon Sep 17 00:00:00 2001
From: Twentysix <Twentysix26@users.noreply.github.com>
Date: Thu, 27 Apr 2017 11:35:15 +0200
Subject: [PATCH] [Core] Handle connection errors, fetch token from env var /
 db

Also reimplemented the old status codes
---
 core/bot.py | 19 +++++++++++++++----
 main.py     | 35 ++++++++++++++++++++++++++++++++---
 2 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/core/bot.py b/core/bot.py
index 1039c6366..05ef984c5 100644
--- a/core/bot.py
+++ b/core/bot.py
@@ -1,12 +1,13 @@
 from discord.ext import commands
 from collections import Counter
 from core.utils.helpers import JsonGuildDB
+from enum import Enum
 import os
 
 
 class Red(commands.Bot):
     def __init__(self, cli_flags, **kwargs):
-        self._shutdown_mode = None
+        self._shutdown_mode = ExitCodes.CRITICAL
         self.db = JsonGuildDB("core/data/settings.json",
                               autosave=True,
                               create_dirs=True)
@@ -48,14 +49,24 @@ class Red(commands.Bot):
             for page in pages:
                 await ctx.send(page)
 
-    async def logout(self, *, restart=False):
+    async def shutdown(self, *, restart=False):
         """Gracefully quits Red with exit code 0
 
         If restart is True, the exit code will be 26 instead
         Upon receiving that exit code, the launcher restarts Red"""
-        self._shutdown_mode = not restart
-        await super().logout()
+        if not restart:
+            self._shutdown_mode = ExitCodes.SHUTDOWN
+        else:
+            self._shutdown_mode = ExitCodes.RESTART
+
+        await self.logout()
 
     def list_packages(self):
         """Lists packages present in the cogs the folder"""
         return os.listdir("cogs")
+
+
+class ExitCodes(Enum):
+    CRITICAL = 1
+    SHUTDOWN = 0
+    RESTART  = 26
\ No newline at end of file
diff --git a/main.py b/main.py
index a40543a64..fafadb482 100644
--- a/main.py
+++ b/main.py
@@ -1,8 +1,10 @@
-from core.bot import Red
+from core.bot import Red, ExitCodes
 from core.global_checks import init_global_checks
 from core.events import init_events
 from core.json_flusher import init_flusher
 from core.settings import parse_cli_flags
+import asyncio
+import discord
 import logging.handlers
 import logging
 import os
@@ -46,9 +48,11 @@ def init_loggers(cli_flags):
     logger.addHandler(fhandler)
     logger.addHandler(stdout_handler)
 
+    return logger
+
 if __name__ == '__main__':
     cli_flags = parse_cli_flags()
-    init_loggers(cli_flags)
+    log = init_loggers(cli_flags)
     init_flusher()
     description = "Red v3 - Alpha"
     red = Red(cli_flags, description=description, pm_help=None)
@@ -57,4 +61,29 @@ if __name__ == '__main__':
     red.load_extension('core')
     if cli_flags.dev:
         pass # load dev cog here?
-    red.run(os.environ['RED_TOKEN'], bot=not cli_flags.not_bot)
+
+    token = os.environ.get("RED_TOKEN", red.db.get_global("token", None))
+
+    if token is None:
+        log.critical("No token to login with")
+        sys.exit(1)
+
+    loop = asyncio.get_event_loop()
+
+    try:
+        loop.run_until_complete(red.start(token, bot=not cli_flags.not_bot))
+    except discord.LoginFailure:
+        log.critical("This token doesn't seem to be valid. If it belongs to "
+                     "a user account, remember that the --not-bot flag "
+                     "must be used. For self-bot functionalities instead, "
+                     "--self-bot")
+    except KeyboardInterrupt:
+        log.info("Keyboard interrupt detected. Quitting...")
+        loop.run_until_complete(red.logout())
+        red._shutdown_mode = ExitCodes.SHUTDOWN
+    except Exception as e:
+        log.critical("Fatal exception", exc_info=e)
+        loop.run_until_complete(red.logout())
+    finally:
+        sys.exit(red._shutdown_mode.value)
+