From d227591521c2938e48428a8d3933f0cba8706d74 Mon Sep 17 00:00:00 2001
From: Neozonz <neozonz@gmail.com>
Date: Mon, 20 Jan 2014 11:16:28 -0800
Subject: [PATCH] [FEATURE] Create workers as they connect into pool_worker,
 validate usernames before creation

---
 mining/DBInterface.py | 16 +++++++++++++---
 mining/DB_Mysql.py    | 19 +++++++++++++++++++
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/mining/DBInterface.py b/mining/DBInterface.py
index 3bf802d..c28907f 100644
--- a/mining/DBInterface.py
+++ b/mining/DBInterface.py
@@ -4,6 +4,7 @@ from datetime import datetime
 import Queue
 import signal
 import Cache
+from sets import Set
 
 import lib.settings as settings
 
@@ -154,6 +155,13 @@ class DBInterface():
         if username == "":
             log.info("Rejected worker for blank username")
             return False
+        allowed_chars = Set('0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ_-.')
+        if Set(username).issubset(allowed_chars) != True:
+            log.info("Username contains bad arguments")
+            return False
+        if username.count('.') > 1:
+            log.info("Username contains multiple . ")
+            return False
         
         # Force username and password to be strings
         username = str(username)
@@ -166,9 +174,11 @@ class DBInterface():
             self.cache.set(username, password)
             return True
         elif settings.USERS_AUTOADD == True:
-            self.insert_user(username, password)
-            self.cache.set(username, password)
-            return True
+			if self.dbi.get_uid(username) != False:
+				uid = self.dbi.get_uid(username)
+				self.dbi.insert_worker(uid, username, password)
+                self.cache.set(username, password)
+				return True
         
         log.info("Authentication for %s failed" % username)
         return False
diff --git a/mining/DB_Mysql.py b/mining/DB_Mysql.py
index 5db4592..531a1f9 100644
--- a/mining/DB_Mysql.py
+++ b/mining/DB_Mysql.py
@@ -204,6 +204,25 @@ class DB_Mysql():
         
         user = self.dbc.fetchone()
         return user
+		
+    def get_uid(self, id_or_username):
+        log.debug("Finding user id of %s", id_or_username)
+        uname = id_or_username.split(".", 1)[0]
+        self.execute("SELECT `id` FROM `accounts` where username = %s", (uname))
+        row = self.dbc.fetchone()
+
+        if row is None:
+            return False
+        else:
+            uid = row[0]
+            return uid
+
+    def insert_worker(self, account_id, username, password):
+        log.debug("Adding new worker %s", username)
+        query = "INSERT INTO pool_worker"
+        self.execute(query + '(account_id, username, password) VALUES (%s, %s, %s);', (account_id, username, password))
+        self.dbh.commit()
+        return str(username)
         
 
     def delete_user(self, id_or_username):