add more encryption methods

config.json

@@ -3,5 +3,6 @@
     "server_port":8388,
     "local_port":1080,
     "password":"barfoo!",
-    "timeout":600
+    "timeout":600,
+    "method":null
 }

encrypt.py

@@ -0,0 +1,171 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2012 clowwindy
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import sys
+import hashlib
+import string
+import struct
+import logging
+
+
+def random_string(length):
+    import M2Crypto.Rand
+    return M2Crypto.Rand.rand_bytes(length)
+
+
+def get_table(key):
+    m = hashlib.md5()
+    m.update(key)
+    s = m.digest()
+    (a, b) = struct.unpack('<QQ', s)
+    table = [c for c in string.maketrans('', '')]
+    for i in xrange(1, 1024):
+        table.sort(lambda x, y: int(a % (ord(x) + i) - a % (ord(y) + i)))
+    return table
+
+encrypt_table = None
+decrypt_table = None
+
+
+def init_table(key, method=None):
+    if method == 'table':
+        method = None
+    if method:
+        try:
+            __import__('M2Crypto')
+        except ImportError:
+            logging.error('M2Crypto is required to use encryption other than default method')
+            sys.exit(1)
+    if not method:
+        global encrypt_table, decrypt_table
+        encrypt_table = ''.join(get_table(key))
+        decrypt_table = string.maketrans(encrypt_table, string.maketrans('', ''))
+    else:
+        try:
+            Encryptor(key, method)  # make an Encryptor to test if the settings if OK
+        except Exception as e:
+            logging.error(e)
+            sys.exit(1)
+
+
+def EVP_BytesToKey(password, key_len, iv_len):
+    # equivalent to OpenSSL's EVP_BytesToKey() with count 1
+    # so that we make the same key and iv as nodejs version
+    # TODO: cache the results
+    m = []
+    i = 0
+    while len(''.join(m)) < (key_len + iv_len):
+        md5 = hashlib.md5()
+        data = password
+        if i > 0:
+            data = m[i - 1] + password
+        md5.update(data)
+        m.append(md5.digest())
+        i += 1
+    ms = ''.join(m)
+    key = ms[:key_len]
+    iv = ms[key_len:key_len + iv_len]
+    return (key, iv)
+
+
+method_supported = {
+    'aes-128-cfb': (16, 16),
+    'aes-192-cfb': (24, 16),
+    'aes-256-cfb': (32, 16),
+    'bf-cfb': (16, 8),
+    'camellia-128-cfb': (16, 16),
+    'camellia-192-cfb': (24, 16),
+    'camellia-256-cfb': (32, 16),
+    'cast5-cfb': (16, 8),
+    'des-cfb': (8, 8),
+    'idea-cfb': (16, 8),
+    'rc2-cfb': (8, 8),
+    'rc4': (16, 0),
+    'seed-cfb': (16, 16),
+}
+
+
+class Encryptor(object):
+    def __init__(self, key, method=None):
+        if method == 'table':
+            method = None
+        self.key = key
+        self.method = method
+        self.iv = None
+        self.iv_sent = False
+        self.cipher_iv = ''
+        self.decipher = None
+        if method is not None:
+            self.cipher = self.get_cipher(key, method, 1, iv=random_string(32))
+        else:
+            self.cipher = None
+
+    def get_cipher_len(self, method):
+        method = method.lower()
+        m = method_supported.get(method, None)
+        return m
+
+    def iv_len(self):
+        return len(self.cipher_iv)
+
+    def get_cipher(self, password, method, op, iv=None):
+        import M2Crypto.EVP
+        password = password.encode('utf-8')
+        method = method.lower()
+        m = self.get_cipher_len(method)
+        if m:
+            key, iv_ = EVP_BytesToKey(password, m[0], m[1])
+            if iv is None:
+                iv = iv_[:m[1]]
+            if op == 1:
+                self.cipher_iv = iv[:m[1]]  # this iv is for cipher, not decipher
+            return M2Crypto.EVP.Cipher(method.replace('-', '_'), key, iv, op, key_as_bytes=0, d='md5', salt=None, i=1, padding=1)
+
+        logging.error('method %s not supported' % method)
+        sys.exit(1)
+
+    def encrypt(self, buf):
+        if len(buf) == 0:
+            return buf
+        if self.method is None:
+            return string.translate(buf, encrypt_table)
+        else:
+            if self.iv_sent:
+                return self.cipher.update(buf)
+            else:
+                self.iv_sent = True
+                return self.cipher_iv + self.cipher.update(buf)
+
+    def decrypt(self, buf):
+        if len(buf) == 0:
+            return buf
+        if self.method is None:
+            return string.translate(buf, decrypt_table)
+        else:
+            if self.decipher is None:
+                decipher_iv_len = self.get_cipher_len(self.method)[1]
+                decipher_iv = buf[:decipher_iv_len]
+                self.decipher = self.get_cipher(self.key, self.method, 0, iv=decipher_iv)
+                buf = buf[decipher_iv_len:]
+                if len(buf) == 0:
+                    return buf
+            return self.decipher.update(buf)

local.py

@@ -27,31 +27,11 @@ if sys.version_info < (2, 6):
 else:
     import json
 import struct
-import string
-import hashlib
 import os
 import logging
 import getopt
 import socket
-
+import encrypt
-
-def get_table(key):
-    m = hashlib.md5()
-    m.update(key)
-    s = m.digest()
-    (a, b) = struct.unpack('<QQ', s)
-    table = [c for c in string.maketrans('', '')]
-    for i in xrange(1, 1024):
-        table.sort(lambda x, y: int(a % (ord(x) + i) - a % (ord(y) + i)))
-    return table
-
-
-def encrypt(data):
-    return data.translate(encrypt_table)
-
-
-def decrypt(data):
-    return data.translate(decrypt_table)
 
 
 class RemoteHandler(object):
@@ -65,14 +45,14 @@ class RemoteHandler(object):
         conn.connect((SERVER, REMOTE_PORT))
 
     def on_connect(self, s):
-        self.conn.write(encrypt(self.local_handler.addr_to_send))
+        self.conn.write(self.local_handler.encryptor.encrypt(self.local_handler.addr_to_send))
         for piece in self.local_handler.cached_pieces:
-            self.conn.write(encrypt(piece))
+            self.conn.write(self.local_handler.encryptor.encrypt(piece))
         # TODO write cached pieces
         self.local_handler.stage = 5
 
     def on_data(self, s, data):
-        data = decrypt(data)
+        data = self.local_handler.encryptor.decrypt(data)
         self.local_handler.conn.write(data)
 
     def on_close(self, s):
@@ -86,7 +66,7 @@ class RemoteHandler(object):
 class LocalHandler(object):
     def on_data(self, s, data):
         if self.stage == 5:
-            data = encrypt(data)
+            data = self.encryptor.encrypt(data)
             self.remote_handler.conn.write(data)
             return
         if self.stage == 0:
@@ -152,6 +132,7 @@ class LocalHandler(object):
         self.addr_to_send = ''
         self.conn = conn
         self.cached_pieces = []
+        self.encryptor = encrypt.Encryptor(KEY, METHOD)
 
         conn.on('data', self.on_data)
         conn.on('end', self.on_end)
@@ -173,6 +154,7 @@ if __name__ == '__main__':
     REMOTE_PORT = config['server_port']
     PORT = config['local_port']
     KEY = config['password']
+    METHOD = config.get('method', None)
 
     argv = sys.argv[1:]
     if '-6' in sys.argv[1:]:
@@ -180,7 +162,7 @@ if __name__ == '__main__':
 
     level = logging.INFO
 
-    optlist, args = getopt.getopt(argv, 's:p:k:l:v')
+    optlist, args = getopt.getopt(argv, 's:p:k:l:m:v')
     for key, value in optlist:
         if key == '-p':
             REMOTE_PORT = int(value)
@@ -190,14 +172,15 @@ if __name__ == '__main__':
             PORT = int(value)
         elif key == '-s':
             SERVER = value
+        elif key == '-m':
+            METHOD = value
         elif key == '-v':
             level = logging.NOTSET
 
     logging.basicConfig(level=level, format='%(asctime)s %(levelname)1.1s %(message)s',
                         datefmt='%Y-%m-%d %H:%M:%S', filemode='a+')
 
-    encrypt_table = ''.join(get_table(KEY))
+    encrypt.init_table(KEY, METHOD)
-    decrypt_table = string.maketrans(encrypt_table, string.maketrans('', ''))
     try:
         logging.info("starting server at port %d ..." % PORT)
         loop = ssloop.instance()

server.py

@@ -27,31 +27,11 @@ if sys.version_info < (2, 6):
 else:
     import json
 import struct
-import string
-import hashlib
 import os
 import logging
 import getopt
 import socket
-
+import encrypt
-
-def get_table(key):
-    m = hashlib.md5()
-    m.update(key)
-    s = m.digest()
-    (a, b) = struct.unpack('<QQ', s)
-    table = [c for c in string.maketrans('', '')]
-    for i in xrange(1, 1024):
-        table.sort(lambda x, y: int(a % (ord(x) + i) - a % (ord(y) + i)))
-    return table
-
-
-def encrypt(data):
-    return data.translate(encrypt_table)
-
-
-def decrypt(data):
-    return data.translate(decrypt_table)
 
 
 class RemoteHandler(object):
@@ -72,7 +52,7 @@ class RemoteHandler(object):
         self.local_handler.stage = 5
 
     def on_data(self, s, data):
-        data = encrypt(data)
+        data = self.local_handler.encryptor.encrypt(data)
         self.local_handler.conn.write(data)
 
     def on_close(self, s):
@@ -85,7 +65,7 @@ class RemoteHandler(object):
 
 class LocalHandler(object):
     def on_data(self, s, data):
-        data = decrypt(data)
+        data = self.encryptor.decrypt(data)
         if self.stage == 5:
             self.remote_handler.conn.write(data)
             return
@@ -143,6 +123,7 @@ class LocalHandler(object):
         self.addr_to_send = ''
         self.conn = conn
         self.cached_pieces = []
+        self.encryptor = encrypt.Encryptor(KEY, METHOD)
 
         conn.on('data', self.on_data)
         conn.on('end', self.on_end)
@@ -162,6 +143,7 @@ if __name__ == '__main__':
         config = json.load(f)
     PORT = config['server_port']
     KEY = config['password']
+    METHOD = config.get('method', None)
 
     argv = sys.argv[1:]
     if '-6' in sys.argv[1:]:
@@ -169,20 +151,21 @@ if __name__ == '__main__':
 
     level = logging.INFO
 
-    optlist, args = getopt.getopt(argv, 'p:k:v')
+    optlist, args = getopt.getopt(sys.argv[1:], 'p:k:m:')
     for key, value in optlist:
         if key == '-p':
             PORT = int(value)
         elif key == '-k':
             KEY = value
+        elif key == '-m':
+            METHOD = value
         elif key == '-v':
             level = logging.NOTSET
 
     logging.basicConfig(level=level, format='%(asctime)s %(levelname)1.1s %(message)s',
                         datefmt='%Y-%m-%d %H:%M:%S', filemode='a+')
 
-    encrypt_table = ''.join(get_table(KEY))
+    encrypt.init_table(KEY, METHOD)
-    decrypt_table = string.maketrans(encrypt_table, string.maketrans('', ''))
     try:
         logging.info("starting server at port %d ..." % PORT)
         loop = ssloop.instance()