Merge pull request #93 from tvincentNuoDB/master

mjmichaels · mjmichaels · commit 979c80d4674e · 2015-06-04T14:29:36.000-04:00
More Python 3 support and beginning of unencrypted session handling
diff --git a/pynuodb/connection.py b/pynuodb/connection.py
@@ -15,14 +15,10 @@
 from .util import getCloudEntry
 
 import time
-import string
-
-# http://www.python.org/dev/peps/pep-0249
 
 apilevel = "2.0"
 threadsafety = 1
 paramstyle = "qmark"
-#schema='user', auto_commit=False
 
 def connect(database, host, user, password, options=None):
     """Creates a connection object.
@@ -97,10 +93,12 @@ def __init__(self, dbName, broker, username, password, options):
         parameters = {'user' : username, 'timezone' : time.strftime('%Z')}
         if options:
             parameters.update(options)
+            if 'cipher' in options and options['cipher'] == 'None':
+                self.__session.set_encryption(False)
 
         version, serverKey, salt = self.__session.open_database(dbName, parameters, cp)
-
-        sessionKey = cp.computeSessionKey(string.upper(username), password, salt, serverKey)
+            
+        sessionKey = cp.computeSessionKey(username.upper(), password, salt, serverKey)
         self.__session.setCiphers(RC4Cipher(sessionKey), RC4Cipher(sessionKey))
 
         self.__session.check_auth()
diff --git a/pynuodb/crypt.py b/pynuodb/crypt.py
@@ -21,16 +21,22 @@
 
 import hashlib
 import random
-import string
 import binascii
+import sys
+
+systemVersion = sys.version[0]
 
 def toHex(bigInt):
-    hexStr = (hex(bigInt)[2:])[:-1]
+    #Python 3 will no longer insert an L for type formatting
+    if systemVersion is '3':
+        hexStr = (hex(bigInt)[2:])
+    else:
+        hexStr = (hex(bigInt)[2:])[:-1]
     # if the number is the right size then the hex string will be missing one
     # character that some platforms assume, so force an even length encoding
     if len(hexStr) % 2 == 1:
         hexStr = "0" + hexStr
-    return string.upper(hexStr)
+    return hexStr.upper()
 
 def fromHex(hexStr):
     return int(hexStr, 16)
@@ -87,6 +93,9 @@ def toByteString(bigInt):
 def fromByteString(byteStr):
     result = 0
     shiftCount = 0
+    if systemVersion == '3':
+        if type(byteStr) is bytes:
+            byteStr = byteStr.decode('latin-1')
     for b in reversed(byteStr):
         result = result | ((ord(b) & 0xff) << shiftCount)
         shiftCount = shiftCount + 8
@@ -111,11 +120,17 @@ def __init__(self, primeStr=defaultPrime, generatorStr=defaultGenerator):
         primeBytes = toByteString(self.__primeInt)
         generatorBytes = toByteString(self.__generatorInt)
         paddingLength = len(primeBytes) - len(generatorBytes)
+        paddingBuffer = chr(0) * paddingLength
 
         md = hashlib.sha1()
+        if systemVersion == '3':
+            primeBytes = primeBytes.encode('latin-1') 
+            generatorBytes = generatorBytes.encode('latin-1') 
+            paddingBuffer = paddingBuffer.encode('latin-1') 
+
         md.update(primeBytes)
         if paddingLength > 0:
-            md.update(chr(0) * paddingLength)
+            md.update(paddingBuffer)
         md.update(generatorBytes)
 
         self.__k = fromByteString(md.digest())
@@ -134,9 +149,13 @@ class RemotePassword:
     def __init__(self):
         self.__group = RemoteGroup()
 
+
     def _getUserHash(self, account, password, salt):
         md = hashlib.sha1()
-        md.update(account + ":" + password)
+        userInfo = account + ":" + password
+        if systemVersion == '3':
+            userInfo = userInfo.encode('latin-1')
+        md.update(userInfo)
         hash1 = md.digest()
 
         md = hashlib.sha1()
@@ -150,6 +169,11 @@ def _computeScramble(self, clientPublicKey, serverPublicKey):
         serverBytes = toByteString(serverPublicKey)
 
         md = hashlib.sha1()
+
+        if systemVersion == '3':
+            clientBytes = clientBytes.encode('latin-1') 
+            serverBytes = serverBytes.encode('latin-1') 
+
         md.update(clientBytes)
         md.update(serverBytes)
 
@@ -186,6 +210,8 @@ def computeSessionKey(self, account, password, salt, serverKey):
         secretBytes = toByteString(sessionSecret)
 
         md = hashlib.sha1()
+        if systemVersion == '3':
+            secretBytes = secretBytes.encode('latin-1')
         md.update(secretBytes)
 
         return md.digest()
@@ -235,27 +261,42 @@ def computeSessionKey(self, clientKey, verifier):
 class RC4Cipher:
 
     def __init__(self, key):
-        self.__S = range(256)
-        self.__s1 = 0
-        self.__s2 = 0
+        if systemVersion == '3':
+            self.__state = list(range(256))
+            key = key.decode('latin-1')
+        else:
+            self.__state = range(256)
+        self.__idx1 = 0
+        self.__idx2 = 0
 
-        state = self.__S
+        state = self.__state
 
         j = 0
         for i in range(256):
-            j = (j + state[i] + ord(key[i % len(key)])) % 256
+            byteString = key[i % len(key)]
+            byteString = ord(byteString)
+
+            j = (j + state[i] + byteString) % 256
             state[i], state[j] = state[j], state[i]
 
     def transform(self, data):
         transformed = []
-        state = self.__S
-
+        state = self.__state
+        if type(data) is bytes:
+            data.decode("latin-1")
         for char in data:
-            self.__s1 = (self.__s1 + 1) % 256
-            self.__s2 = (self.__s2 + state[self.__s1]) % 256
-            state[self.__s1], state[self.__s2] = state[self.__s2], state[self.__s1]
-            cipherByte = ord(char) ^ state[(state[self.__s1] + state[self.__s2]) % 256]
-            transformed.append(chr(cipherByte))
 
+            self.__idx1 = (self.__idx1 + 1) % 256
+            self.__idx2 = (self.__idx2 + state[self.__idx1]) % 256
+            state[self.__idx1], state[self.__idx2] = state[self.__idx2], state[self.__idx1]
+            cipherByte = ord(char) ^ state[(state[self.__idx1] + state[self.__idx2]) % 256]
+            transformed.append(chr(cipherByte))
         return ''.join(transformed)
 
+class NoCipher:
+
+    def __init__(self):
+        pass
+
+    def transform(self, data):
+        return data
diff --git a/pynuodb/encodedsession.py b/pynuodb/encodedsession.py
@@ -10,7 +10,7 @@
 import struct
 import decimal
 
-from .crypt import toByteString, fromByteString, toSignedByteString, fromSignedByteString
+from .crypt import toByteString, fromByteString, toSignedByteString, fromSignedByteString, NoCipher
 from .session import Session, SessionException
 from . import protocol
 from . import datatype
@@ -59,7 +59,7 @@ class EncodedSession(Session):
                 supporting function.
     exchangeMessages -- Exchange the pending message for an optional response from the server.
     setCiphers -- Re-sets the incoming and outgoing ciphers for the session.
-    
+    set_encryption -- Takes a value of type boolean. Setting encryption to False will result in disabling encryption after the handshake
     Private Functions:
     __init__ -- Constructor for the EncodedSession class.
     _peekTypeCode -- Looks at the next Type Code off the session. (Does not move inpos)
@@ -80,6 +80,8 @@ def __init__(self, host, port, service='SQL2'):
         self.__inpos = 0
         """ @type : int """
         self.closed = False
+        """ @type : boolean """
+        self.__encryption = True
 
     # Mostly for connections
     def open_database(self, db_name, parameters, cp):
@@ -91,20 +93,21 @@ def open_database(self, db_name, parameters, cp):
         self._putMessageId(protocol.OPENDATABASE).putInt(protocol.CURRENT_PROTOCOL_VERSION).putString(db_name).putInt(len(parameters))
         for (k, v) in parameters.items():
             self.putString(k).putString(v)
+
         self.putNull().putString(cp.genClientKey())
 
         self._exchangeMessages()
-
         version = self.getInt()
         serverKey = self.getString()
         salt = self.getString()
-
         return version, serverKey, salt
 
     def check_auth(self):
         try:
             self._putMessageId(protocol.AUTHENTICATION).putString(protocol.AUTH_TEST_STR)
             self._exchangeMessages()
+            if self.__encryption is False:
+                self._setCiphers(NoCipher(), NoCipher())
         except SessionException as e:
             raise ProgrammingError('Failed to authenticate: ' + str(e))
 
@@ -136,6 +139,9 @@ def send_rollback(self):
         self._putMessageId(protocol.ROLLBACKTRANSACTION)
         self._exchangeMessages()
 
+    def set_encryption(self, value):
+        self.__encryption = value
+
     def test_connection(self):
         # Create a statement handle
         self._putMessageId(protocol.CREATE)
@@ -776,7 +782,7 @@ def getValue(self):
     def _exchangeMessages(self, getResponse=True):
         """Exchange the pending message for an optional response from the server."""
         try:
-            # print "message to server: %s" %  (self.__output)
+            #print("message to server: %s" %  (self.__output))
             self.send(self.__output)
         finally:
             self.__output = None
diff --git a/pynuodb/protocol.py b/pynuodb/protocol.py
@@ -6,9 +6,6 @@
 NULL                            = 1
 TRUE                            = 2
 FALSE                           = 3
-UUID                            = 200
-SCALEDCOUNT1                    = 199
-SCALEDCOUNT2                    = 225
 INTMINUS10                      = 10
 INTMINUS1                       = 19
 INT0                            = 20
@@ -17,33 +14,35 @@
 INTLEN8                         = 59
 SCALEDLEN0                      = 60
 SCALEDLEN8                      = 68
-DOUBLELEN0                      = 77
-DOUBLELEN8                      = 85
 UTF8COUNT1                      = 69
 UTF8COUNT4                      = 72
-UTF8LEN0                        = 109
-UTF8LEN39                       = 148
 OPAQUECOUNT1                    = 73
 OPAQUECOUNT4                    = 76
-OPAQUELEN0                      = 149
-OPAQUELEN39                     = 188
-BLOBLEN0                        = 189
-BLOBLEN4                        = 193
-CLOBLEN0                        = 194
-CLOBLEN4                        = 198
+DOUBLELEN0                      = 77
+DOUBLELEN8                      = 85
 MILLISECLEN0                    = 86     # milliseconds since January 1, 1970
 MILLISECLEN8                    = 94
 NANOSECLEN0                     = 95     # nanoseconds since January 1, 1970
 NANOSECLEN8                     = 103
 TIMELEN0                        = 104    # milliseconds since midnight
 TIMELEN4                        = 108
+UTF8LEN0                        = 109
+UTF8LEN39                       = 148
+OPAQUELEN0                      = 149
+OPAQUELEN39                     = 188
+BLOBLEN0                        = 189
+BLOBLEN4                        = 193
+CLOBLEN0                        = 194
+CLOBLEN4                        = 198
+SCALEDCOUNT1                    = 199
+UUID                            = 200
+SCALEDDATELEN1                  = 201
+SCALEDDATELEN8                  = 208
 SCALEDTIMELEN1                  = 209
 SCALEDTIMELEN8                  = 216
 SCALEDTIMESTAMPLEN1             = 217
 SCALEDTIMESTAMPLEN8             = 224
-SCALEDDATELEN1                  = 201
-SCALEDDATELEN8                  = 208
-
+SCALEDCOUNT2                    = 225
 
 # Protocol Messages
 FAILURE                         = 0
@@ -126,8 +125,8 @@
 ATTACHDEBUGGER                  = 77 
 DEBUGREQUEST                    = 78 
 GETSEQUENCEVALUE2               = 79 
-GETLIMIT                        = 80 
-SETLIMIT                        = 81 
+GETCONNECTIONLIMIT              = 80 
+SETCONNECTIONLIMIT              = 81 
 DELETEBLOBDATA                  = 82 
 EXECUTEBATCH                    = 83
 EXECUTEBATCHPREPAREDSTATEMENT   = 84
@@ -149,6 +148,27 @@
 SUPPORTSTRANSACTIONISOLATION    = 100 
 GETCATALOG                      = 101 
 GETCURRENTSCHEMA                = 102 
+PREPARECALL                     = 103
+EXECUTECALLABLESTATEMENT        = 104 
+SETQUERYTIMEOUT                 = 105 
+GETPROCEDURES                   = 106 
+GETPROCEDURECOLUMNS             = 107 
+GETSUPERTABLES                  = 108
+GETSUPERTYPES                   = 109 
+GETFUNCTIONS                    = 110 
+GETFUNCTIONCOLUMNS              = 111 
+GETTABLEPRIVILEGES              = 112 
+GETCOLUMNPRIVILEGES             = 113 
+GETCROSSREFERENCE               = 114 
+ALLPROCEDURESARECALLABLE        = 115 
+ALLTABLESARESELECTABLE          = 116 
+GETATTRIBUTES                   = 117 
+GETUDTS                         = 118
+GETVERSIONCOLUMNS               = 119 
+GETLOBCHUNK                     = 120
+GETLASTSTATEMENTTIMEMICROS      = 121
+
+
 
 
 # Error code values
@@ -320,6 +340,10 @@ def lookup_code(error_code):
 PROTOCOL_VERSION13         = 13  # 01/24/2013   Added some JDBC methods
 PROTOCOL_VERSION14         = 14  # 02/18/2014   Changed strings in Types/TypeUtil.cpp - affected Python driver
 PROTOCOL_VERSION15         = 15  # 04/22/2014   openDatabase returns server-side connection id in response
+PROTOCOL_VERSION16         = 16  # 07/01/2014
+PROTOCOL_VERSION17         = 17  # 11/03/2014 Support for client sending last commit info
+PROTOCOL_VERSION18         = 18  # 02/19/2015 JDBC metadata updates
+PROTOCOL_VERSION19         = 19  # 05/01/2015 Server timing of statements
 #
 # The current protocol version of THIS driver.  The server will negotiate the lowest compatible version.
 #
diff --git a/pynuodb/session.py b/pynuodb/session.py
diff --git a/tests/nuodb_crypt_test.py b/tests/nuodb_crypt_test.py
