From fde194f9b091f6f03ee82ebfd82b800c5f79bb6f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Peter=20=C5=A0urda?= <surda@economicsofbitcoin.com>
Date: Mon, 19 Feb 2018 16:57:47 +0100
Subject: [PATCH] More messagetypes hardening

- only load modules at program start
---
 src/messagetypes/__init__.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/messagetypes/__init__.py b/src/messagetypes/__init__.py
index 1a5223df..06783eac 100644
--- a/src/messagetypes/__init__.py
+++ b/src/messagetypes/__init__.py
@@ -3,6 +3,7 @@ from os import path, listdir
 from string import lower
 
 from debug import logger
+import messagetypes
 import paths
 
 class MsgBase(object):
@@ -15,9 +16,8 @@ def constructObject(data):
     if data[""] not in whitelist:
         return None
     try:
-        m = import_module("messagetypes." + data[""])
-        classBase = getattr(m, data[""].title())
-    except (NameError, ImportError):
+        classBase = getattr(getattr(messagetypes, data[""]), data[""].title())
+    except (NameError, AttributeError):
         logger.error("Don't know how to handle message type: \"%s\"", data[""], exc_info=True)
         return None
     try:
@@ -43,7 +43,7 @@ else:
         if splitted[1] != ".py":
             continue
         try:
-            import_module("." + splitted[0], "messagetypes")
+            import_module(".{}".format(splitted[0]), "messagetypes")
         except ImportError:
             logger.error("Error importing %s", mod, exc_info=True)
         else: