From 74d902d14f1ed3b262f5e497f5e4ec3c0a3c6b99 Mon Sep 17 00:00:00 2001
From: Rohitt Vashishtha <aero31aero@gmail.com>
Date: Wed, 28 Aug 2019 04:49:29 +0530
Subject: [PATCH] botserver: Allow importing custom bot modules.

We can now specify path to a bot's python file as the ini section
header in the botserver's config file. For example:

[~/Documents/helloworld.py]
email=a@b.com
key=XXXX
site=https://b.com
token=XXXX
---
 zulip_botserver/tests/test_server.py      | 32 +++++++++++++++++++++++
 zulip_botserver/zulip_botserver/server.py | 17 ++++++++++--
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/zulip_botserver/tests/test_server.py b/zulip_botserver/tests/test_server.py
index 57329c3..eb1601c 100644
--- a/zulip_botserver/tests/test_server.py
+++ b/zulip_botserver/tests/test_server.py
@@ -4,6 +4,8 @@ from typing import Any, Dict
 import unittest
 from .server_test_lib import BotServerTestCase
 import json
+from importlib import import_module
+from types import ModuleType
 
 from zulip_botserver import server
 from zulip_botserver.input_parameters import parse_args
@@ -177,5 +179,35 @@ class BotServerTests(BotServerTestCase):
         }
         assert json.dumps(bot_conf2, sort_keys=True) == json.dumps(expected_config2, sort_keys=True)
 
+    def test_load_lib_modules(self) -> None:
+        # This testcase requires hardcoded paths, which here is a good thing so if we ever
+        # restructure zulip_bots, this test would fail and we would also update Botserver
+        # at the same time.
+        helloworld = import_module('zulip_bots.bots.{bot}.{bot}'.format(bot='helloworld'))
+        root_dir = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../../'))
+        # load valid module name
+        module = server.load_lib_modules(['helloworld'])['helloworld']
+        assert module == helloworld
+
+        # load valid file path
+        path = os.path.join(root_dir, 'zulip_bots/zulip_bots/bots/{bot}/{bot}.py'.format(bot='helloworld'))
+        module = server.load_lib_modules([path])[path]
+        assert module.__name__ == 'custom_bot_module'
+        assert module.__file__ == path
+        assert isinstance(module, ModuleType)
+
+        # load invalid module name
+        with self.assertRaisesRegexp(SystemExit,  # type: ignore
+                                     'Error: Bot "botserver-test-case-random-bot" doesn\'t exist. '
+                                     'Please make sure you have set up the botserverrc file correctly.'):
+            module = server.load_lib_modules(['botserver-test-case-random-bot'])['botserver-test-case-random-bot']
+
+        # load invalid file path
+        with self.assertRaisesRegexp(SystemExit,  # type: ignore
+                                     'Error: Bot "{}/zulip_bots/zulip_bots/bots/helloworld.py" doesn\'t exist. '
+                                     'Please make sure you have set up the botserverrc file correctly.'.format(root_dir)):
+            path = os.path.join(root_dir, 'zulip_bots/zulip_bots/bots/{bot}.py'.format(bot='helloworld'))
+            module = server.load_lib_modules([path])[path]
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/zulip_botserver/zulip_botserver/server.py b/zulip_botserver/zulip_botserver/server.py
index 69ecfdb..00a0110 100644
--- a/zulip_botserver/zulip_botserver/server.py
+++ b/zulip_botserver/zulip_botserver/server.py
@@ -5,11 +5,13 @@ import logging
 import json
 import os
 import sys
+import importlib.util
 
 from configparser import MissingSectionHeaderError, NoOptionError
 from flask import Flask, request
 from importlib import import_module
 from typing import Any, Dict, Union, List, Optional
+from types import ModuleType
 from werkzeug.exceptions import BadRequest, Unauthorized
 
 from zulip import Client
@@ -68,13 +70,24 @@ def parse_config_file(config_file_path: str) -> configparser.ConfigParser:
     parser.read(config_file_path)
     return parser
 
+def load_module_from_file(file_path: str) -> ModuleType:
+    # Wrapper around importutil; see https://stackoverflow.com/a/67692/3909240.
+    spec = importlib.util.spec_from_file_location("custom_bot_module", file_path)
+    lib_module = importlib.util.module_from_spec(spec)
+    assert spec is not None
+    assert spec.loader is not None
+    spec.loader.exec_module(lib_module)
+    return lib_module
 
 def load_lib_modules(available_bots: List[str]) -> Dict[str, Any]:
     bots_lib_module = {}
     for bot in available_bots:
         try:
-            module_name = 'zulip_bots.bots.{bot}.{bot}'.format(bot=bot)
-            lib_module = import_module(module_name)
+            if bot.endswith('.py') and os.path.isfile(bot):
+                lib_module = load_module_from_file(bot)
+            else:
+                module_name = 'zulip_bots.bots.{bot}.{bot}'.format(bot=bot)
+                lib_module = import_module(module_name)
             bots_lib_module[bot] = lib_module
         except ImportError:
             error_message = ("Error: Bot \"{}\" doesn't exist. Please make sure "