extract find_library

shadowsocks/crypto/ctypes_libsodium.py

@@ -23,10 +23,11 @@
 from __future__ import absolute_import, division, print_function, \
     with_statement
 
-import logging
-from ctypes import CDLL, c_char_p, c_int, c_ulonglong, byref, \
+from ctypes import c_char_p, c_int, c_ulonglong, byref, \
     create_string_buffer, c_void_p
 
+from shadowsocks.crypto import util
+
 __all__ = ['ciphers']
 
 libsodium = None
@@ -41,21 +42,11 @@ BLOCK_SIZE = 64
 def load_libsodium():
     global loaded, libsodium, buf
 
-    from ctypes.util import find_library
-    libsodium_path = None
-    for p in ('sodium', 'libsodium'):
-        libsodium_path = find_library(p)
-        if libsodium_path:
-            break
-    else:
-        import glob
-        for libsodium_path in glob.glob('/usr/lib/libsodium.*'):
-            pass
-    if libsodium_path is None:
+    libsodium = util.find_library('sodium', 'crypto_stream_salsa20_xor_ic',
+                                  'libsodium')
+    if libsodium is None:
         raise Exception('libsodium not found')
-    logging.info('loading libsodium from %s', libsodium_path)
-    libsodium = CDLL(libsodium_path)
-    libsodium.sodium_init.restype = c_int
+
     libsodium.crypto_stream_salsa20_xor_ic.restype = c_int
     libsodium.crypto_stream_salsa20_xor_ic.argtypes = (c_void_p, c_char_p,
                                                        c_ulonglong,
@@ -67,8 +58,6 @@ def load_libsodium():
                                                         c_char_p, c_ulonglong,
                                                         c_char_p)
 
-    libsodium.sodium_init()
-
     buf = create_string_buffer(buf_size)
     loaded = True
 
@@ -118,8 +107,6 @@ ciphers = {
 
 
 def test_salsa20():
-    from shadowsocks.crypto import util
-
     cipher = Salsa20Crypto(b'salsa20', b'k' * 32, b'i' * 16, 1)
     decipher = Salsa20Crypto(b'salsa20', b'k' * 32, b'i' * 16, 0)
 
@@ -127,7 +114,6 @@ def test_salsa20():
 
 
 def test_chacha20():
-    from shadowsocks.crypto import util
 
     cipher = Salsa20Crypto(b'chacha20', b'k' * 32, b'i' * 16, 1)
     decipher = Salsa20Crypto(b'chacha20', b'k' * 32, b'i' * 16, 0)

shadowsocks/crypto/ctypes_openssl.py

@@ -23,10 +23,11 @@
 from __future__ import absolute_import, division, print_function, \
     with_statement
 
-import logging
-from ctypes import CDLL, c_char_p, c_int, c_long, byref,\
+from ctypes import c_char_p, c_int, c_long, byref,\
     create_string_buffer, c_void_p
 
+from shadowsocks.crypto import util
+
 __all__ = ['ciphers']
 
 libcrypto = None
@@ -38,40 +39,12 @@ buf_size = 2048
 def load_openssl():
     global loaded, libcrypto, buf
 
-    from ctypes.util import find_library
-    libcrypto_path = None
-    for p in ('crypto', 'eay32', 'libeay32'):
-        libcrypto_path = find_library(p)
-        if libcrypto_path:
-            break
-    else:
-        # We may get here when find_library fails because, for example,
-        # the user does not have sufficient privileges to access those
-        # tools underlying find_library on linux.
-
-        import glob
-        import sys
-
-        patterns = ['/usr/lib/libcrypto.*']
-
-        # Some linux distros may store so in alternative locations
-        if sys.maxsize > 2 ** 32:
-            # Python is 64-bit
-            patterns.extend(['/usr/lib64/libcrypto.*'])
-        else:
-            # Python is 32-bit
-            patterns.extend(['/usr/lib32/libcrypto.*'])
-
-        for pat in patterns:
-            files = glob.glob(pat)
-            if files:
-                libcrypto_path = files[0]
-                break
-
-    if libcrypto_path is None:
+    libcrypto = util.find_library(('crypto', 'eay32'),
+                                  'EVP_get_cipherbyname',
+                                  'libcrypto')
+    if libcrypto is None:
         raise Exception('libcrypto(OpenSSL) not found')
-    logging.info('loading libcrypto from %s', libcrypto_path)
-    libcrypto = CDLL(libcrypto_path)
+
     libcrypto.EVP_get_cipherbyname.restype = c_void_p
     libcrypto.EVP_CIPHER_CTX_new.restype = c_void_p
 
@@ -173,7 +146,6 @@ ciphers = {
 
 
 def run_method(method):
-    from shadowsocks.crypto import util
 
     cipher = CtypesCrypto(method, b'k' * 32, b'i' * 16, 1)
     decipher = CtypesCrypto(method, b'k' * 32, b'i' * 16, 0)

shadowsocks/crypto/util.py

@@ -20,6 +20,57 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+from __future__ import absolute_import, division, print_function, \
+    with_statement
+
+import logging
+
+
+def find_library(possible_lib_names, search_symbol, library_name):
+    from ctypes.util import find_library
+    from ctypes import CDLL
+
+    paths = []
+
+    if type(possible_lib_names) is not list:
+        possible_lib_names = [possible_lib_names]
+
+    for name in possible_lib_names:
+        path = find_library(name)
+        if path:
+            paths.append(path)
+
+    if not paths:
+        # We may get here when find_library fails because, for example,
+        # the user does not have sufficient privileges to access those
+        # tools underlying find_library on linux.
+        import glob
+
+        for name in possible_lib_names:
+            patterns = [
+                '/usr/local/lib*/lib%s.*' % name,
+                '/usr/lib*/lib%s.*' % name,
+                'lib%s.*' % name,
+                '%s.dll' % name,
+                'lib%s.dll' % name]
+
+            for pat in patterns:
+                files = glob.glob(pat)
+                if files:
+                    paths.extend(files)
+    for path in paths:
+        try:
+            lib = CDLL(path)
+            if hasattr(lib, search_symbol):
+                logging.info('loading %s from %s', library_name, path)
+                return lib
+            else:
+                logging.warn('can\'t find symbol %s in %s', search_symbol,
+                             path)
+        except Exception:
+            pass
+    return None
+
 
 def run_cipher(cipher, decipher):
     from os import urandom
@@ -49,3 +100,17 @@ def run_cipher(cipher, decipher):
     end = time.time()
     print('speed: %d bytes/s' % (BLOCK_SIZE * rounds / (end - start)))
     assert b''.join(results) == plain
+
+
+def test_find_library():
+    assert find_library('c', 'strcpy', 'libc') is not None
+    assert find_library(['c'], 'strcpy', 'libc') is not None
+    assert find_library('crypto', 'EVP_CipherUpdate', 'libcrypto') is not None
+    assert find_library('notexist', 'strcpy', 'libnotexist') is None
+    assert find_library('c', 'symbol_not_exist', 'c') is None
+    assert find_library(['notexist', 'c', 'crypto'],
+                        'EVP_CipherUpdate', 'libc') is not None
+
+
+if __name__ == '__main__':
+    test_find_library()