Global refactoring

8 files changed, 84 insertions, 30 deletions

diff --git a/lib/__init__.py b/lib/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/lib/__init__.py
diff --git a/lib/config.py b/lib/config.py
new file mode 100644
index 0000000..f4fab70
--- /dev/null
+++ b/lib/config.py
@@ -0,0 +1,16 @@
+import configparser
+
+class Config:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if not cls._instance:
+            cls._instance = super(Config, cls).__new__(
+                cls, *args, **kwargs)
+
+        return cls._instance
+
+    def __init__(self):
+        self.parser = configparser.ConfigParser()
+        self.parser.read('pywhoisd.conf') # Find a way not to hardcode this
+
diff --git a/core.py b/lib/core.py
index 300f6ff..e7ac571 100644
--- a/core.py
+++ b/lib/core.py
@@ -2,10 +2,14 @@ import socketserver
 import ipaddr
 import re
 
+from lib.config import Config
+
 class Daemon():
     def __init__(self, data):
         self.data = data
         self.domain_regexp = re.compile("(?!-)[A-Z\d-]{1,63}(?<!-)$", re.IGNORECASE)
+        self.footer = None
+        self.header = None
 
     def query(self, q):
         if self.is_ip(q):
@@ -87,13 +91,21 @@ class Daemon():
     def get_help(self):
         return "This will be the help"
 
-    # TODO
     def get_footer(self):
-        return "This will be the footer"
+        if not self.footer:
+            f = open(Config().parser['Printing']['footer'])
+            self.footer = f.read()
+            f.close()
+
+        return self.footer
 
-    # TODO
     def get_header(self):
-        return "This will be the header"
+        if not self.header:
+            f = open(Config().parser['Printing']['header'])
+            self.header = f.read()
+            f.close()
+
+        return self.header
 
 class WhoisHandler(socketserver.BaseRequestHandler):
 
@@ -104,24 +116,24 @@ class WhoisHandler(socketserver.BaseRequestHandler):
         data = str(self.request.recv(100).strip(), 'utf-8')
         print('Received: {}'.format(data))
 
-        #response = self.daemon.get_header() + " \n"
-        response = self.daemon.query(data)
-        #response += self.daemon.get_footer() + "\n"
+        response = self.daemon.get_header()
+        response += self.daemon.query(data)
+        response += self.daemon.get_footer()
         
         self.request.sendall(bytes(response, 'utf-8'))
     
 class ClassicServer(socketserver.ThreadingTCPServer):    
-    def __init__(self, config, daemon):
-        host = config['Servers']['classic_host']
-        port = int(config['Servers']['classic_port'])
+    def __init__(self, daemon):
+        host = Config().parser['Servers']['classic_host']
+        port = int(Config().parser['Servers']['classic_port'])
         self.daemon = daemon
 
         socketserver.ThreadingTCPServer.__init__(self, (host, port), WhoisHandler)
 
 class WebServer(socketserver.ThreadingTCPServer):
-    def __init__(self, config, daemon):
-        self.host = config['Servers']['web_host']
-        self.port = int(config['Servers']['web_port'])
+    def __init__(self, daemon):
+        self.host = Config().parser['Servers']['web_host']
+        self.port = int(Config().parser['Servers']['web_port'])
         self.daemon = daemon
 
         socketserver.ThreadingTCPServer.__init__(self, (self.host, self.port), WhoisHandler)
diff --git a/model.py b/lib/model.py
index ef5ce28..933d901 100644
--- a/model.py
+++ b/lib/model.py
@@ -1,5 +1,7 @@
 from xml.etree.ElementTree import ElementTree
 
+from lib.config import Config
+
 class Network():
     """A simple network definition"""
     
@@ -58,8 +60,8 @@ class Person():
 class Data():
     """Abstract class for storing and getting information"""
     
-    def __init__(self, config):
-        self.config = config
+    def __init__(self):
+        self.config = Config().parser
         self.networks = None
         self.domains = None
         self.persons = None
diff --git a/pywhoisd.conf b/pywhoisd.conf
index b18da39..d01cfed 100644
--- a/pywhoisd.conf
+++ b/pywhoisd.conf
@@ -6,7 +6,7 @@ classic = yes
 
 # Only makes sense when classic server is enabled
 classic_host = localhost
-classic_port = 4344
+classic_port = 4343
 
 # Run a web whois server?
 web = no
@@ -20,4 +20,8 @@ web_port = 8080
 mode = xml
 
 # Only makes sense when xml storage mode is enabled
-xml_file = examples/networks.xml
\ No newline at end of file
+xml_file = examples/networks.xml
+
+[Printing]
+header = res/header.txt
+footer = res/footer.txt
diff --git a/pywhoisd.py b/pywhoisd.py
index 11ad133..ec793e7 100755
--- a/pywhoisd.py
+++ b/pywhoisd.py
@@ -1,16 +1,18 @@
 #!/usr/bin/python3
 import configparser
 import concurrent.futures
+import signal
+import sys
 
-import core
-import model
+from lib import core
+from lib import model
+from lib.config import Config
 
 class PyWhoisD():
     """Main class. It reads the configuration options and starts the server"""
 
     def __init__(self):
-        self.config = configparser.ConfigParser()
-        self.config.read('pywhoisd.conf')
+        self.config = Config().parser
 
         self.data = None
         self.daemon = None
@@ -19,17 +21,20 @@ class PyWhoisD():
 
         self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=2)
 
-    # What kind of storage are we using?
+    def signal_sigint(self, signal, frame):
+        print("[+] Received SIGINT signal. Aborting...")
+
+        sys.exit(0)
+
+   # What kind of storage are we using?
     def config_data(self):
-        """Config data sources.
+        """Config data sources."""
 
-        At the moment only XML is supported.
-        """
-        
+        #At the moment only XML is supported.
         mode = self.config['Storage']['mode']
 
         if mode == 'xml':
-            self.data = model.DataXML(self.config)
+            self.data = model.DataXML()
 
     def config_daemon(self):
         """Config common information source for all configured servers"""
@@ -53,12 +58,12 @@ class PyWhoisD():
         """Sets up server configuration from config files"""
         
         if self.classic():
-            self.classic_server = core.ClassicServer(self.config, self.daemon)            
+            self.classic_server = core.ClassicServer(self.daemon)            
         else:
             print("[+] Classic server is not enabled")
                 
         if self.web():
-            self.web_server = core.WebServer(self.config, self.daemon)
+            self.web_server = core.WebServer(self.daemon)
         else:
             print("[+] Web server is not enabled")
 
@@ -77,11 +82,13 @@ class PyWhoisD():
 
 
     def main(self):
+        signal.signal(signal.SIGINT, self.signal_sigint)
+
         self.config_daemon()
         self.start_servers()
 
         # Wait for running server to finish. Probably never.
-        self.executor.shutdown()
+        # self.executor.shutdown()
 
 if __name__ == "__main__":
     pwd = PyWhoisD()
diff --git a/res/footer.txt b/res/footer.txt
new file mode 100644
index 0000000..6139265
--- /dev/null
+++ b/res/footer.txt
@@ -0,0 +1,2 @@
+
+-----END WHOIS ANSWER-----
diff --git a/res/header.txt b/res/header.txt
new file mode 100644
index 0000000..5eaf530
--- /dev/null
+++ b/res/header.txt
@@ -0,0 +1,11 @@
+           @@@@@    @@@@@  @@@@@    @@@@@  @@@@@         @@@@@@@@@@@@          
+           @@@@@    @@@@@  @@@@@@   @@@@@  @@@@@         @@@@@@@@@@@@@          
+           @@@@    @@@@    @@@@@@  @@@@    @@@@          @@@@    @@@@           
+          @@@@    @@@@    @@@@@@@ @@@@    @@@@          @@@@@@@@@@@@            
+         @@@@    @@@@    @@@@@@@@@@@@    @@@@          @@@@@@@@@@@
+        @@@@    @@@@    @@@@ @@@@@@@    @@@@          @@@@
+       @@@@@@@@@@@@   @@@@@  @@@@@@@  @@@@@@@@@@@@  @@@@@
+       @@@@@@@@@@    @@@@@   @@@@@@  @@@@@@@@@@@@  @@@@@
+
+-----BEGIN WHOIS ANSWER-----
+