Detect firmware version and choose correct driver automatically

Change-Id: I8c26fb85af2e85ee93439a618c0b428282a32883
diff --git a/lavapdu/__init__.py b/lavapdu/__init__.py
index 9b73b54..6c6aaf6 100644
--- a/lavapdu/__init__.py
+++ b/lavapdu/__init__.py
@@ -19,4 +19,4 @@
 #  MA 02110-1301, USA.
 
 import pdurunner
-import socketserver
+import socketserver
\ No newline at end of file
diff --git a/lavapdu/apcdrivers.py b/lavapdu/apcdrivers.py
deleted file mode 100644
index 0dad9b6..0000000
--- a/lavapdu/apcdrivers.py
+++ /dev/null
@@ -1,168 +0,0 @@
-#! /usr/bin/python
-
-#  Copyright 2013 Linaro Limited
-#  Author Matt Hart <matthew.hart@linaro.org>
-#
-#  This program is free software; you can redistribute it and/or modify
-#  it under the terms of the GNU General Public License as published by
-#  the Free Software Foundation; either version 2 of the License, or
-#  (at your option) any later version.
-#
-#  This program is distributed in the hope that it will be useful,
-#  but WITHOUT ANY WARRANTY; without even the implied warranty of
-#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-#  GNU General Public License for more details.
-#
-#  You should have received a copy of the GNU General Public License
-#  along with this program; if not, write to the Free Software
-#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
-#  MA 02110-1301, USA.
-
-import logging
-
-from driver import PDUDriver
-
-
-class apc7952(PDUDriver):
-    def _pdu_logout(self):
-        self._back_to_main()
-        logging.debug("Logging out")
-        self.connection.send("4\r")
-
-    def _back_to_main(self):
-        logging.debug("Returning to main menu")
-        self.connection.expect('>')
-        for i in range(1, 20):
-            #print("Sending escape character")
-            self.connection.send("\x1B")
-            self.connection.send("\r")
-            res = self.connection.expect(["4- Logout", "> "])
-            if res == 0:
-                logging.debug("Back at main menu")
-                break
-        #self.connection.send("\r")
-        #self.connection.expect("4- Logout")
-        #self.connection.expect("> ")
-
-    def _enter_outlet(self, outlet, enter_needed=True):
-        outlet = "%s" % outlet
-        logging.debug("Attempting to enter outlet %s", outlet)
-        if (enter_needed):
-            self.connection.expect("Press <ENTER> to continue...")
-        logging.debug("Sending enter")
-        self.connection.send("\r")
-        self.connection.expect("> ")
-        logging.debug("Sending outlet number")
-        self.connection.send(outlet)
-        self.connection.send("\r")
-        logging.debug("Finished entering outlet")
-
-    def _port_interaction(self, command, port_number):
-        print("Attempting command: %s port: %i" % (command, port_number))
-        ### make sure in main menu here
-        self._back_to_main()
-        self.connection.send("\r")
-        self.connection.expect("1- Device Manager")
-        self.connection.expect("> ")
-        logging.debug("Entering Device Manager")
-        self.connection.send("1\r")
-        res = self.connection.expect(["3- Outlet Control/Configuration", "2- Outlet Control", "2- Outlet Management", "------- Device Manager"])
-        logging.debug("Matched pattern %s", res)
-        if res == 0:
-            self.connection.send("3\r")
-            self._enter_outlet(port_number)
-        elif res == 1:
-            self.connection.send("2\r")
-            self._enter_outlet(port_number)
-        elif res == 2:
-            self.connection.send("2\r")
-        elif res == 3:
-            logging.debug("Matched ------- Device Manager")
-            self._enter_outlet(port_number, False)
-        res = self.connection.expect(["1- Control Outlet", "1- Outlet Control/Configuration"])
-        self.connection.expect("> ")
-        self.connection.send("1\r")
-        res = self.connection.expect(["> ", "Press <ENTER> to continue..."])
-        if res == 1:
-            logging.debug("Stupid paging thingmy detected, pressing enter")
-            self.connection.send("\r")
-        self.connection.send("\r")
-        res = self.connection.expect(["Control Outlet %s" % port_number, "Control Outlet"])
-        if res == 0:
-            logging.debug("Already at the right port")
-        else:
-            self.connection.send("%s\r" % port_number)
-            self.connection.send("1\r")
-        self.connection.expect("3- Immediate Reboot")
-        self.connection.expect("> ")
-        if command == "reboot":
-            self.connection.send("3\r")
-            self.connection.expect("Immediate Reboot")
-            self._do_it()
-        elif command == "delayed":
-            self.connection.send("6\r")
-            self.connection.expect("Delayed Reboot")
-            self._do_it()
-        elif command == "on":
-            self.connection.send("1\r")
-            self.connection.expect("Immediate On")
-            self._do_it()
-        elif command == "off":
-            self.connection.send("2\r")
-            self.connection.expect("Immediate Off")
-            self._do_it()
-        else:
-            logging.debug("Unknown command!")
-
-    def _do_it(self):
-        self.connection.expect("Enter 'YES' to continue or <ENTER> to cancel :")
-        self.connection.send("YES\r")
-        self.connection.expect("Press <ENTER> to continue...")
-        self.connection.send("\r")
-
-    def port_delayed(self, port_number):
-        self._port_interaction("delayed", port_number)
-
-    def port_on(self, port_number):
-        self._port_interaction("on", port_number)
-
-    def port_off(self, port_number):
-        self._port_interaction("off", port_number)
-
-    def port_reboot(self, port_number):
-        self._port_interaction("reboot", port_number)
-
-
-class apc8959(PDUDriver):
-    connection = None
-    pdu_commands = {"off": "olOff", "on": "olOn", "reboot": "olReboot", "delayed": "olDlyReboot"}
-
-    def _pdu_logout(self):
-        logging.debug("logging out")
-        self.connection.send("\r")
-        self.connection.send("exit")
-        self.connection.send("\r")
-        logging.debug("done")
-
-    def _pdu_get_to_prompt(self):
-        self.connection.send("\r")
-        self.connection.expect('apc>')
-
-    def _port_interaction(self, command, port_number):
-        logging.debug("Attempting %s on port %i" % (command, port_number))
-        self._pdu_get_to_prompt()
-        self.connection.sendline(self.pdu_commands[command] + (" %i" % port_number))
-        self.connection.expect("E000: Success")
-        logging.debug("done")
-
-    def port_delayed(self, port_number):
-        self._port_interaction("delayed", port_number)
-
-    def port_on(self, port_number):
-        self._port_interaction("on", port_number)
-
-    def port_off(self, port_number):
-        self._port_interaction("off", port_number)
-
-    def port_reboot(self, port_number):
-        self._port_interaction("reboot", port_number)
diff --git a/lavapdu/dbhandler.py b/lavapdu/dbhandler.py
new file mode 100644
index 0000000..bc89847
--- /dev/null
+++ b/lavapdu/dbhandler.py
@@ -0,0 +1,60 @@
+#! /usr/bin/python
+
+#  Copyright 2013 Linaro Limited
+#  Author Matt Hart <matthew.hart@linaro.org>
+#
+#  This program is free software; you can redistribute it and/or modify
+#  it under the terms of the GNU General Public License as published by
+#  the Free Software Foundation; either version 2 of the License, or
+#  (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU General Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License
+#  along with this program; if not, write to the Free Software
+#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
+#  MA 02110-1301, USA.
+
+import logging
+import psycopg2
+import time
+
+class DBHandler(object):
+    def __init__(self, config):
+        logging.debug("Creating new DBHandler: %s" % config["dbhost"])
+        logging.getLogger().name = "DBHandler"
+        self.conn = psycopg2.connect(database=config["dbname"], user=config["dbuser"],
+                                     password=config["dbpass"], host=config["dbhost"])
+        self.cursor = self.conn.cursor()
+
+    def do_sql(self, sql):
+        logging.debug("executing sql: %s" % sql)
+        self.cursor.execute(sql)
+        self.conn.commit()
+
+    def do_sql_with_fetch(self, sql):
+        logging.debug("executing sql: %s" % sql)
+        self.cursor.execute(sql)
+        row = self.cursor.fetchone()
+        self.conn.commit()
+        return row
+
+    def delete_row(self, row_id):
+        logging.debug("deleting row %i" % row_id)
+        self.do_sql("delete from pdu_queue where id=%i" % row_id)
+
+    def get_res(self, sql):
+        return self.cursor.execute(sql)
+
+    def get_next_job(self):
+        now = int(time.time())
+        row = self.do_sql_with_fetch("select id,hostname,port,request from pdu_queue where exectime < %i order by id asc limit 1" % now)
+        return row
+
+    def close(self):
+        logging.debug("Closing DBHandler")
+        self.cursor.close()
+        self.conn.close()
\ No newline at end of file
diff --git a/lavapdu/drivers/__init__.py b/lavapdu/drivers/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/lavapdu/drivers/__init__.py
diff --git a/lavapdu/drivers/apc7952.py b/lavapdu/drivers/apc7952.py
new file mode 100644
index 0000000..447bac1
--- /dev/null
+++ b/lavapdu/drivers/apc7952.py
@@ -0,0 +1,95 @@
+#! /usr/bin/python
+
+#  Copyright 2013 Linaro Limited
+#  Author Matt Hart <matthew.hart@linaro.org>
+#
+#  This program is free software; you can redistribute it and/or modify
+#  it under the terms of the GNU General Public License as published by
+#  the Free Software Foundation; either version 2 of the License, or
+#  (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU General Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License
+#  along with this program; if not, write to the Free Software
+#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
+#  MA 02110-1301, USA.
+
+import logging
+from driver import PDUDriver
+
+class apc7952(PDUDriver):
+
+    def _pdu_logout(self):
+        self._back_to_main()
+        logging.debug("Logging out")
+        self.connection.send("4\r")
+
+    def _back_to_main(self):
+        logging.debug("Returning to main menu")
+        self.connection.expect('>')
+        for i in range(1, 20):
+            self.connection.send("\x1B")
+            self.connection.send("\r")
+            res = self.connection.expect(["4- Logout", "> "])
+            if res == 0:
+                logging.debug("Back at main menu")
+                break
+
+    def _enter_outlet(self, outlet, enter_needed=True):
+        outlet = "%s" % outlet
+        logging.debug("Attempting to enter outlet %s", outlet)
+        if (enter_needed):
+            self.connection.expect("Press <ENTER> to continue...")
+        logging.debug("Sending enter")
+        self.connection.send("\r")
+        self.connection.expect("> ")
+        logging.debug("Sending outlet number")
+        self.connection.send(outlet)
+        self.connection.send("\r")
+        logging.debug("Finished entering outlet")
+
+    def _port_interaction(self, command, port_number):
+        print("Attempting command: %s port: %i" % (command, port_number))
+        ### make sure in main menu here
+        self._back_to_main()
+        self.connection.send("\r")
+        self.connection.expect("1- Device Manager")
+        self.connection.expect("> ")
+        logging.debug("Entering Device Manager")
+        self.connection.send("1\r")
+        self.connection.expect("------- Device Manager")
+        self.connection.send("2\r")
+        res = self.connection.expect("1- Outlet Control/Configuration")
+        self.connection.expect("> ")
+        self.connection.send("1\r")
+        self._enter_outlet(port_number, False)
+        self.connection.expect("> ")
+        self.connection.send("1\r")
+        res = self.connection.expect(["> ", "Press <ENTER> to continue..."])
+        if res == 1:
+            logging.debug("Stupid paging thingmy detected, pressing enter")
+            self.connection.send("\r")
+        self.connection.send("\r")
+        if command == "on":
+            self.connection.send("1\r")
+            self.connection.expect("Immediate On")
+            self._do_it()
+        elif command == "off":
+            self.connection.send("2\r")
+            self.connection.expect("Immediate Off")
+            self._do_it()
+        else:
+            logging.debug("Unknown command!")
+
+    def _do_it(self):
+        self.connection.expect("Enter 'YES' to continue or <ENTER> to cancel :")
+        self.connection.send("YES\r")
+        self.connection.expect("Press <ENTER> to continue...")
+        self.connection.send("\r")
+
+    class Meta():
+        handled_firmware = ["v3.7.3"]
\ No newline at end of file
diff --git a/lavapdu/drivers/apc8959.py b/lavapdu/drivers/apc8959.py
new file mode 100644
index 0000000..a3f1e29
--- /dev/null
+++ b/lavapdu/drivers/apc8959.py
@@ -0,0 +1,47 @@
+#! /usr/bin/python
+
+#  Copyright 2013 Linaro Limited
+#  Author Matt Hart <matthew.hart@linaro.org>
+#
+#  This program is free software; you can redistribute it and/or modify
+#  it under the terms of the GNU General Public License as published by
+#  the Free Software Foundation; either version 2 of the License, or
+#  (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU General Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License
+#  along with this program; if not, write to the Free Software
+#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
+#  MA 02110-1301, USA.
+
+import logging
+from driver import PDUDriver
+
+
+class apc8959(PDUDriver):
+    pdu_commands = {"off": "olOff", "on": "olOn"}
+
+    def _pdu_logout(self):
+        logging.debug("logging out")
+        self.connection.send("\r")
+        self.connection.send("exit")
+        self.connection.send("\r")
+        logging.debug("done")
+
+    def _pdu_get_to_prompt(self):
+        self.connection.send("\r")
+        self.connection.expect('apc>')
+
+    def _port_interaction(self, command, port_number):
+        logging.debug("Attempting %s on port %i" % (command, port_number))
+        self._pdu_get_to_prompt()
+        self.connection.sendline(self.pdu_commands[command] + (" %i" % port_number))
+        self.connection.expect("E000: Success")
+        logging.debug("done")
+
+    class Meta():
+        handled_firmware = ["v5.1.9"]
\ No newline at end of file
diff --git a/lavapdu/drivers/apc9218.py b/lavapdu/drivers/apc9218.py
new file mode 100644
index 0000000..a81f6d7
--- /dev/null
+++ b/lavapdu/drivers/apc9218.py
@@ -0,0 +1,62 @@
+#! /usr/bin/python
+
+#  Copyright 2013 Linaro Limited
+#  Author Matt Hart <matthew.hart@linaro.org>
+#
+#  This program is free software; you can redistribute it and/or modify
+#  it under the terms of the GNU General Public License as published by
+#  the Free Software Foundation; either version 2 of the License, or
+#  (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU General Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License
+#  along with this program; if not, write to the Free Software
+#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
+#  MA 02110-1301, USA.
+
+import logging
+from apc7952 import apc7952
+
+class apc9218(apc7952):
+    handled_firmware = ["v3.0.3","v3.0.1"]
+
+    def _port_interaction(self, command, port_number):
+        print("Attempting command: %s port: %i" % (command, port_number))
+        ### make sure in main menu here
+        self._back_to_main()
+        self.connection.send("\r")
+        self.connection.expect("1- Device Manager")
+        self.connection.expect("> ")
+        logging.debug("Entering Device Manager")
+        self.connection.send("1\r")
+        self.connection.expect("------- Device Manager")
+        logging.debug("Got to Device Manager")
+        self._enter_outlet(port_number, False)
+        res = self.connection.expect(["1- Control Outlet", "1- Outlet Control/Configuration"])
+        self.connection.expect("> ")
+        self.connection.send("1\r")
+        res = self.connection.expect(["> ", "Press <ENTER> to continue..."])
+        if res == 1:
+            logging.debug("Stupid paging thingmy detected, pressing enter")
+            self.connection.send("\r")
+        self.connection.send("\r")
+        res = self.connection.expect(["Control Outlet %s" % port_number, "Control Outlet"])
+        self.connection.expect("3- Immediate Reboot")
+        self.connection.expect("> ")
+        if command == "on":
+            self.connection.send("1\r")
+            self.connection.expect("Immediate On")
+            self._do_it()
+        elif command == "off":
+            self.connection.send("2\r")
+            self.connection.expect("Immediate Off")
+            self._do_it()
+        else:
+            logging.debug("Unknown command!")
+
+    class Meta():
+        handled_firmware = ["v3.0.3","v3.0.1"]
\ No newline at end of file
diff --git a/lavapdu/driver.py b/lavapdu/drivers/driver.py
similarity index 66%
rename from lavapdu/driver.py
rename to lavapdu/drivers/driver.py
index bd36c18..7be55e9 100644
--- a/lavapdu/driver.py
+++ b/lavapdu/drivers/driver.py
@@ -18,10 +18,29 @@
 #  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 #  MA 02110-1301, USA.
 
+import logging
 
-class PDUDriver():
+
+class NoDriverException(Exception):
+    pass
+
+
+class PDUDriver(object):
     connection = None
-    pdu_commands = {"off": "olOff", "on": "olOn", "reboot": "olReboot", "delayed": "olDlyReboot"}
+    firmware_dict = {}
 
     def __init__(self, connection):
         self.connection = connection
+
+        # return the driver that provides the firmware version
+    def _port_interaction(self, command, port_number):
+        pass
+
+    def _pdu_logout(self):
+        pass
+
+    def port_on(self, port_number):
+        self._port_interaction("on", port_number)
+
+    def port_off(self, port_number):
+        self._port_interaction("off", port_number)
\ No newline at end of file
diff --git a/lavapdu/engine.py b/lavapdu/engine.py
index beb954f..7bcbfa7 100644
--- a/lavapdu/engine.py
+++ b/lavapdu/engine.py
@@ -21,35 +21,53 @@
 import pexpect
 import os
 import logging
+import pkgutil
 import sys
 
-from apcdrivers import apc8959
-from apcdrivers import apc7952
-
 
 class PDUEngine():
     connection = None
     prompt = 0
     driver = None
+    firmware_dict = {}
 
     def __init__(self, pdu_hostname, pdu_telnetport=23):
         self.exec_string = "/usr/bin/telnet %s %d" % (pdu_hostname, pdu_telnetport)
         logging.debug("Created new PDUEngine: %s" % self.exec_string)
-        #self.connection.logfile_read = sys.stdout
-        prompt = self._pdu_login("apc", "apc")
-        if prompt == 0:
-            logging.debug("Found a v5 prompt")
-            self.driver = apc8959(self.connection)
-        elif prompt == 1:
-            logging.debug("Found a v3 prompt")
-            self.driver = apc7952(self.connection)
+        required_version = self._pdu_login("apc", "apc")
+        logging.debug("Got firmware version: %s" % required_version)
+        driver_list = self.load_all_modules_from_dir("drivers")
+        for driver in driver_list:
+            handled = []
+            exec("handled = %s.Meta.handled_firmware" % driver)
+            for firmware_value in handled:
+                self.firmware_dict[firmware_value] = driver
+
+        logging.debug("Firmware versions supported: %s" % self.firmware_dict)
+        if self.firmware_dict[required_version]:
+            driver = self.firmware_dict[required_version]
+            logging.debug("Using driver %s for version: %s" % (driver, required_version))
+            exec("self.driver = %s(self.connection)" % driver)
         else:
-            logging.debug("Unknown prompt!")
+            self.driver = None
+
+    def load_all_modules_from_dir(self, dirname):
+        module_list = []
+        for importer, package_name, _ in pkgutil.iter_modules([dirname]):
+            full_package_name = '%s.%s' % (dirname, package_name)
+            if full_package_name not in sys.modules and (not full_package_name == "%s.driver" % dirname):
+                import_string = "global %s\nfrom %s import %s" % (package_name,full_package_name,package_name)
+                exec (import_string)
+                logging.debug(import_string)
+            if (not full_package_name == "%s.driver" % dirname):
+                module_list.append(package_name)
+        return module_list
 
     def pduconnect(self):
         self.connection = self.get_connection(self.exec_string)
 
     def pduclose(self):
+        logging.debug("Closing connection: %s" % self.connection)
         self.connection.close(True)
 
     def pdureconnect(self):
@@ -58,7 +76,7 @@
 
     def get_connection(self, exec_string):
         connection = pexpect.spawn(exec_string)
-        connection.logfile = sys.stdout
+        #connection.logfile = sys.stdout
         return connection
 
     def is_busy(self):
@@ -68,7 +86,8 @@
 
     def close(self):
         self.driver._pdu_logout()
-        self.connection.close(True)
+        self.firmware_dict = {}
+        del(self)
 
     def _pdu_login(self, username, password):
         logging.debug("attempting login with username %s, password %s" % (username, password))
@@ -78,34 +97,26 @@
         self.connection.send("apc\r")
         self.connection.expect("Password  :")
         self.connection.send("apc\r")
-        return self.connection.expect(["apc>", ">"])
-
+        output = self.connection.read(250)
+        #print "OUTPUT: %s" % output
+        a = output.split("AOS")[1].split()[0]
+        #print "A: %s" % a
+        b = a.strip()
+        #print "B: %s" % b
+        version = b
+        return version.strip()
 
 if __name__ == "__main__":
-    #pe1 = PDUEngine("pdu15")
-    #pe1.driver.port_off(22)
-    #pe1.driver.port_on(22)
-    #pe1.close()
-    #pe2 = PDUEngine("pdu14")
-    #pe2.driver.port_off(6)
-    #pe2.driver.port_on(6)
-    #pe2.close()
-    #pe3 = PDUEngine("pdu01")
-    #pe3.driver.port_reboot(1)
-    #pe3.driver.port_off(1)
-    #pe3.driver.port_on(1)
-    #pe3.close()
     logging.basicConfig(level=logging.DEBUG)
     logging.getLogger().setLevel(logging.DEBUG)
-    pe4 = PDUEngine("192.168.1.153")
-    pe4.driver.port_reboot(1)
-    pe4.driver.port_reboot(2)
-    pe4.driver.port_reboot(3)
-    pe4.driver.port_reboot(4)
-    pe4.driver.port_reboot(5)
-    pe4.driver.port_reboot(6)
-    pe4.driver.port_reboot(7)
-    pe4.driver.port_reboot(8)
-    #pe4.driver.port_off(8)
-    #pe4.driver.port_on(8)
-    pe4.close()
+    pe = PDUEngine("pdu01")
+    pe.driver.port_off(1)
+    pe.driver.port_on(1)
+    pe.close()
+    pe = PDUEngine("pdu03")
+    pe.driver.port_off(3)
+    pe.driver.port_on(3)
+    pe.close()
+
+    #pe = PDUEngine("pdu16")
+    #pe.close()
\ No newline at end of file
diff --git a/lavapdu/pdurunner.py b/lavapdu/pdurunner.py
index cc2b19d..8a83f12 100644
--- a/lavapdu/pdurunner.py
+++ b/lavapdu/pdurunner.py
@@ -21,7 +21,7 @@
 import logging
 import time
 from engine import PDUEngine
-from socketserver import DBHandler
+from dbhandler import DBHandler
 
 
 class PDURunner():
@@ -36,6 +36,7 @@
         job = db.get_next_job()
         if job:
             job_id, hostname, port, request = job
+            logging.debug(job)
             logging.info("Processing queue item: (%s %s) on hostname: %s" % (request, port, hostname))
             #logging.debug(id, hostname, request, port)
             res = self.do_job(hostname, port, request)
@@ -44,27 +45,29 @@
             logging.debug("Found nothing to do in database")
 
     def do_job(self, hostname, port, request):
-        retries = 5
+        retries = 10
         while retries > 0:
             try:
                 pe = PDUEngine(hostname, 23)
-                if request == "reboot":
-                    pe.driver.port_reboot(port)
-                elif request == "on":
+                if request == "on":
                     pe.driver.port_on(port)
+                    return true
                 elif request == "off":
                     pe.driver.port_off(port)
-                elif request == "delayed":
-                    pe.driver.port_delayed(port)
+                    return true
                 else:
                     logging.debug("Unknown request type: %s" % request)
+                    return false
                 pe.pduclose()
+                del(pe)
                 retries = 0
-            except:
-                logging.warn("Failed to execute job: %s %s %s (attempts left %i)" % (hostname, port, request, retries))
+            except Exception as e:
+                logging.warn("Failed to execute job: %s %s %s (attempts left %i) error was %s" %
+                             (hostname, port, request, retries, e.message))
                 #logging.warn(e)
                 time.sleep(5)
                 retries -= 1
+        return false
 
     def run_me(self):
         logging.info("Starting up the PDURunner")
diff --git a/lavapdu/socketserver.py b/lavapdu/socketserver.py
index 5a6011c..62551bc 100644
--- a/lavapdu/socketserver.py
+++ b/lavapdu/socketserver.py
@@ -19,46 +19,10 @@
 #  MA 02110-1301, USA.
 
 import SocketServer
-import psycopg2
 import logging
 import socket
-
-
-class DBHandler(object):
-    def __init__(self, config):
-        logging.debug("Creating new DBHandler: %s" % config["dbhost"])
-        logging.getLogger().name = "DBHandler"
-        self.conn = psycopg2.connect(database=config["dbname"], user=config["dbuser"],
-                                     password=config["dbpass"], host=config["dbhost"])
-        self.cursor = self.conn.cursor()
-
-    def do_sql(self, sql):
-        logging.debug("executing sql: %s" % sql)
-        self.cursor.execute(sql)
-        self.conn.commit()
-
-    def do_sql_with_fetch(self, sql):
-        logging.debug("executing sql: %s" % sql)
-        self.cursor.execute(sql)
-        row = self.cursor.fetchone()
-        self.conn.commit()
-        return row
-
-    def delete_row(self, row_id):
-        logging.debug("deleting row %i" % row_id)
-        self.do_sql("delete from pdu_queue where id=%i" % row_id)
-
-    def get_res(self, sql):
-        return self.cursor.execute(sql)
-
-    def get_next_job(self):
-        row = self.do_sql_with_fetch("select * from pdu_queue order by id asc limit 1")
-        return row
-
-    def close(self):
-        logging.debug("Closing DBHandler")
-        self.cursor.close()
-        self.conn.close()
+import time
+from dbhandler import DBHandler
 
 
 class ListenerServer(object):
@@ -75,8 +39,15 @@
         del(self.db)
 
     def create_db(self):
-        sql = "create table if not exists pdu_queue (id serial, hostname text, port int, request text)"
+        sql = "create table if not exists pdu_queue (id serial, hostname text, port int, request text, exectime int)"
         self.db.do_sql(sql)
+        sql = "select column_name from information_schema.columns where table_name='pdu_queue'" \
+              "and column_name='exectime'"
+        res = self.db.do_sql_with_fetch(sql)
+        if not res:
+            logging.info("Old db schema discovered, upgrading")
+            sql = "alter table pdu_queue add column exectime int default 1"
+            self.db.do_sql(sql)
 
     def start(self):
         logging.info("Starting the ListenerServer")
@@ -88,21 +59,41 @@
     def insert_request(self, data):
         logging.getLogger().name = "TCPRequestHandler"
         array = data.split(" ")
-        if len(array) != 3:
+        delay = 10
+        custom_delay = False
+        now = int(time.time())
+        if len(array) < 3:
             logging.info("Wrong data size")
             raise Exception("Unexpected data")
+        if len(array) == 4:
+            delay = int(array[3])
+            custom_delay = True
         hostname = array[0]
         port = int(array[1])
         request = array[2]
-        if not (request in ["reboot", "on", "off", "delayed"]):
+        if not (request in ["reboot","on","off"]):
             logging.info("Unknown request: %s" % request)
             raise Exception("Unknown request: %s" % request)
+        if request == "reboot":
+            logging.debug("reboot requested, submitting off/on")
+            self.queue_request(hostname,port,"off",now)
+            self.queue_request(hostname,port,"on",now+delay)
+        else:
+            if custom_delay:
+                logging.debug("using delay as requested")
+                self.queue_request(hostname,port,request,now+delay)
+            else:
+                self.queue_request(hostname,port,request,now)
+
+    def queue_request(self, hostname, port, request, exectime):
         db = DBHandler(self.server.config)
-        sql = "insert into pdu_queue (hostname,port,request) values ('%s',%i,'%s')" % (hostname, port, request)
+        sql = "insert into pdu_queue (hostname,port,request,exectime)" \
+              "values ('%s',%i,'%s')" % (hostname,port,request,exectime)
         db.do_sql(sql)
         db.close()
         del(db)
 
+
     def handle(self):
         logging.getLogger().name = "TCPRequestHandler"
         ip = self.client_address[0]
@@ -112,7 +103,7 @@
             try:
                 request_host = socket.gethostbyaddr(ip)[0]
             except socket.herror as e:
-                logging.debug("Unable to resolve: %s error: %s" % (ip, e))
+                logging.debug("Unable to resolve: %s error: %s" % (ip,e))
                 request_host = ip
             logging.info("Received a request from %s: '%s'" % (request_host, data))
             self.insert_request(data)
@@ -122,7 +113,6 @@
             self.request.sendall("nack\n")
         self.request.close()
 
-
 class TCPServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer):
     allow_reuse_address = True
     daemon_threads = True
@@ -133,11 +123,11 @@
     logging.getLogger().setLevel(logging.DEBUG)
     logging.debug("Executing from __main__")
     starter = {"hostname": "0.0.0.0",
-               "port": 16421,
-               "dbhost": "127.0.0.1",
-               "dbuser": "pdudaemon",
-               "dbpass": "pdudaemon",
-               "dbname": "lavapdu",
+               "port":16421,
+               "dbhost":"127.0.0.1",
+               "dbuser":"pdudaemon",
+               "dbpass":"pdudaemon",
+               "dbname":"lavapdu",
                "logging_level": logging.DEBUG}
     ss = ListenerServer(starter)
-    ss.start()
+    ss.start()
\ No newline at end of file
diff --git a/pduclient b/pduclient
index f9cb734..3347c60 100755
--- a/pduclient
+++ b/pduclient
@@ -1,17 +1,36 @@
 #!/usr/bin/python
 
+#  Copyright 2013 Linaro Limited
+#  Author Matt Hart <matthew.hart@linaro.org>
+#
+#  This program is free software; you can redistribute it and/or modify
+#  it under the terms of the GNU General Public License as published by
+#  the Free Software Foundation; either version 2 of the License, or
+#  (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU General Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License
+#  along with this program; if not, write to the Free Software
+#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
+#  MA 02110-1301, USA.
+
 import socket
 import optparse
 
 if __name__ == '__main__':
     usage = "Usage: %prog --daemon deamonhostname --hostname pduhostname --port pduportnum --command pducommand"
     description = "LAVA PDU daemon client"
-    commands = ["reboot", "on", "off", "delayed"]
+    commands = ["reboot", "on", "off"]
     parser = optparse.OptionParser(usage=usage, description=description)
     parser.add_option("--daemon", dest="pdudaemonhostname", action="store", type="string", help="LAVAPDU Daemon (ex: control)")
     parser.add_option("--hostname", dest="pduhostname", action="store", type="string", help="PDU Hostname (ex: pdu05)")
     parser.add_option("--port", dest="pduportnum", action="store", type="string", help="PDU Portnumber (ex: 04)")
-    parser.add_option("--command", dest="pducommand", action="store", type="string", help="PDU command (ex: reboot|on|off|delayed)")
+    parser.add_option("--command", dest="pducommand", action="store", type="string", help="PDU command (ex: reboot|on|off)")
+    parser.add_option("--delay", dest="pdudelay", action="store", type="int", help="Delay before command runs, or between off/on when rebooting")
     (options, args) = parser.parse_args()
     if (not (options.pdudaemonhostname) or not(options.pduhostname) or not (options.pduportnum) or not (options.pducommand)):
         print("Missing option, try -h for help")
@@ -20,7 +39,10 @@
         print("Unknown pdu command: %s" % options.pducommand)
         exit(1)
     #print(options)
-    string = ("%s %s %s" % (options.pduhostname, options.pduportnum, options.pducommand))
+    if options.pdudelay:
+        string = ("%s %s %s" % (options.pduhostname, options.pduportnum, options.pducommand, options.pdudelay))
+    else:
+        string = ("%s %s %s" % (options.pduhostname, options.pduportnum, options.pducommand))
     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     #sock.setblocking(0)  # optional non-blocking
     reply = ""