96 lines
2.8 KiB
Python
96 lines
2.8 KiB
Python
|
import pymysql
|
||
|
|
import json
|
||
|
|
|
||
|
|
class MaDbHelper:
|
||
|
|
def __init__(self):
|
||
|
|
self.mysqlConn = pymysql.connect(host='localhost',
|
||
|
|
user='root',
|
||
|
|
passwd='MlmAccess',
|
||
|
|
db='MlmAccess',
|
||
|
|
port=4306)
|
||
|
|
|
||
|
|
self.mysqlCur = self.mysqlConn.cursor()
|
||
|
|
|
||
|
|
def _initDb(self):
|
||
|
|
"""
|
||
|
|
Generates required tables and users
|
||
|
|
"""
|
||
|
|
|
||
|
|
"""
|
||
|
|
CREATE TABLE vmq_auth_acl
|
||
|
|
(
|
||
|
|
mountpoint VARCHAR(10) NOT NULL,
|
||
|
|
client_id VARCHAR(128) NOT NULL,
|
||
|
|
username VARCHAR(128) NOT NULL,
|
||
|
|
password VARCHAR(128),
|
||
|
|
publish_acl TEXT,
|
||
|
|
subscribe_acl TEXT,
|
||
|
|
CONSTRAINT vmq_auth_acl_primary_key PRIMARY KEY (mountpoint, client_id, username)
|
||
|
|
)
|
||
|
|
"""
|
||
|
|
|
||
|
|
def addUser(self, username, password, publishAclPatterns, subscribeAclPatterns):
|
||
|
|
|
||
|
|
if self.userExists(username):
|
||
|
|
return False
|
||
|
|
|
||
|
|
query = "INSERT INTO `vmq_auth_acl` (`mountpoint`, `client_id`, `username`, `password`, `publish_acl`, `subscribe_acl`) VALUES (%s, %s, %s, PASSWORD(%s), %s, %s);"
|
||
|
|
|
||
|
|
self.mysqlCur.execute(
|
||
|
|
query,
|
||
|
|
(
|
||
|
|
"",
|
||
|
|
username,
|
||
|
|
username,
|
||
|
|
password,
|
||
|
|
self._convertAclPatternList(publishAclPatterns),
|
||
|
|
self._convertAclPatternList(subscribeAclPatterns)
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
self.mysqlConn.commit()
|
||
|
|
|
||
|
|
return True
|
||
|
|
|
||
|
|
def userExists(self, username):
|
||
|
|
query = "SELECT username FROM `vmq_auth_acl` WHERE username=%s"
|
||
|
|
|
||
|
|
self.mysqlCur.execute(
|
||
|
|
query,
|
||
|
|
(username)
|
||
|
|
)
|
||
|
|
|
||
|
|
result = self.mysqlCur.fetchone()
|
||
|
|
return result is not None
|
||
|
|
|
||
|
|
|
||
|
|
def updateUser(self, username, password, publishAclPatterns, subscribeAclPatterns):
|
||
|
|
if self.userExists(username):
|
||
|
|
return self._updateUser(username, password, publishAclPatterns, subscribeAclPatterns)
|
||
|
|
else:
|
||
|
|
return False
|
||
|
|
|
||
|
|
def _updateUser(self, username, password, publishAclPatterns, subscribeAclPatterns):
|
||
|
|
|
||
|
|
query = """
|
||
|
|
UPDATE `vmq_auth_acl`
|
||
|
|
SET `password`=PASSWORD(%s), `publish_acl`=%s, `subscribe_acl`=%s
|
||
|
|
WHERE `username`=%s;
|
||
|
|
"""
|
||
|
|
|
||
|
|
self.mysqlCur.execute(
|
||
|
|
query,
|
||
|
|
(
|
||
|
|
password,
|
||
|
|
self._convertAclPatternList(publishAclPatterns),
|
||
|
|
self._convertAclPatternList(subscribeAclPatterns),
|
||
|
|
username
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
self.mysqlConn.commit()
|
||
|
|
|
||
|
|
def _convertAclPatternList(self, patternList):
|
||
|
|
patternMapLabda = lambda pattern: {"pattern": pattern}
|
||
|
|
patternList = list(map(patternMapLabda, patternList))
|
||
|
|
return json.dumps(patternList)
|