diff --git a/btcrecover/test/test_passwords.py b/btcrecover/test/test_passwords.py index c995aca7..2fd45c23 100644 --- a/btcrecover/test/test_passwords.py +++ b/btcrecover/test/test_passwords.py @@ -22,6 +22,8 @@ import warnings, os, unittest, pickle, tempfile, shutil, multiprocessing, time, gc, filecmp, sys, hashlib, argparse + +from lib.opencl_brute import opencl if __name__ == '__main__': sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) @@ -3287,6 +3289,80 @@ def __init__(self): self.addTests(tl.loadTestsFromTestCase(Test08BIP39Passwords)) +@unittest.skipUnless(has_any_opencl_devices(), "requires OpenCL and a compatible device") +class TestOpenCLPBKDF2BufferSizing(unittest.TestCase): + """Ensure OpenCL PBKDF2 buffer sizes track password limits.""" + + @classmethod + def setUpClass(cls): + super(TestOpenCLPBKDF2BufferSizing, cls).setUpClass() + if not has_any_opencl_devices(): + raise unittest.SkipTest("requires OpenCL and a compatible device") + + import pyopencl as cl + + device = opencl_devices_list[0] + cls.platform_index = 0 + cls.device_index = 0 + for platform_index, platform in enumerate(cl.get_platforms()): + devices = platform.get_devices() + if device in devices: + cls.platform_index = platform_index + cls.device_index = devices.index(device) + break + + cls.algos = opencl.opencl_algos( + cls.platform_index, 0, False, openclDevice=cls.device_index + ) + + def test_sha1_fast_kernel_limits_password_bytes(self): + password = b"a" * 32 + salt = b"b" * 16 + ctx = self.algos.cl_pbkdf2_init( + "sha1", saltlen=len(salt), dklen=16, max_password_bytes=len(password) + ) + buf_structs = ctx[1] + + self.assertEqual(buf_structs.pwdBufferSize_bytes, len(password)) + self.assertEqual(buf_structs.inBufferSize_bytes, len(password)) + + result = self.algos.cl_pbkdf2(ctx, [password], salt, 1000, 16) + expected = [hashlib.pbkdf2_hmac("sha1", password, salt, 1000, 16)] + self.assertEqual(result, expected) + + def test_sha1_falls_back_for_long_passwords(self): + password = b"l" * 200 + salt = b"s" * 16 + ctx = self.algos.cl_pbkdf2_init( + "sha1", saltlen=len(salt), dklen=20, max_password_bytes=len(password) + ) + buf_structs = ctx[1] + + self.assertEqual(buf_structs.pwdBufferSize_bytes, len(password)) + self.assertEqual(buf_structs.inBufferSize_bytes, len(password)) + + result = self.algos.cl_pbkdf2(ctx, [password], salt, 2000, 20) + expected = [hashlib.pbkdf2_hmac("sha1", password, salt, 2000, 20)] + self.assertEqual(result, expected) + + def test_saltlist_buffers_scale_with_password_length(self): + password = b"p" * 144 + salts = [b"salt-one", b"salt-two-long"] + ctx = self.algos.cl_pbkdf2_saltlist_init( + "sha256", pwdlen=len(password), dklen=32 + ) + buf_structs = ctx[1] + + self.assertEqual(buf_structs.pwdBufferSize_bytes, len(password)) + self.assertEqual(buf_structs.inBufferSize_bytes, len(password)) + + result = self.algos.cl_pbkdf2_saltlist(ctx, password, salts, 4096, 32) + expected = [ + hashlib.pbkdf2_hmac("sha256", password, salt, 4096, 32) for salt in salts + ] + self.assertEqual(result, expected) + + if __name__ == '__main__': import argparse diff --git a/lib/opencl_brute/opencl.py b/lib/opencl_brute/opencl.py index 8342169d..d398302a 100644 --- a/lib/opencl_brute/opencl.py +++ b/lib/opencl_brute/opencl.py @@ -8,7 +8,13 @@ from collections import deque from itertools import chain, repeat, zip_longest import numpy as np -import pyopencl as cl + +try: + import pyopencl as cl +except ImportError as _pyopencl_import_error: # pragma: no cover - exercised in environments without PyOpenCL + cl = None +else: + _pyopencl_import_error = None # Minimum number of items to execute in a single OpenCL batch. Some # OpenCL runtimes (such as PoCL) crash when a kernel is launched with a @@ -17,6 +23,15 @@ from lib.opencl_brute.buffer_structs import buffer_structs import os, sys, inspect + +def _require_pyopencl(): + """Ensure PyOpenCL is available before executing OpenCL operations.""" + + if cl is None: + raise ImportError( + "pyopencl is required for OpenCL acceleration but is not installed" + ) from _pyopencl_import_error + current_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) parent_dir = os.path.dirname(current_dir) @@ -67,6 +82,7 @@ def __init__( N_value=15, openclDevice=0, ): + _require_pyopencl() self.workgroupsize = 0 self.computeunits = 0 self.wordSize = None @@ -921,30 +937,70 @@ def func(s, pwdim, pass_g, salt_g, result_g): result = [hexRes[:dklen] for hexRes in result] return result - def cl_pbkdf2_init(self, rtype, saltlen, dklen): + def cl_pbkdf2_init(self, rtype, saltlen, dklen, max_password_bytes=256): bufStructs = buffer_structs() + if max_password_bytes is None: + max_password_bytes = 256 + assert max_password_bytes > 0, "max_password_bytes must be positive" if rtype == "md5": - self.max_out_bytes = bufStructs.specifyMD5(128, saltlen, dklen) + max_in_bytes = max(128, max_password_bytes) + self.max_out_bytes = bufStructs.specifyMD5( + max_in_bytes, + saltlen, + dklen, + max_password_bytes=max_password_bytes, + ) # hmac is defined in with pbkdf2, as a kernel function prg = self.opencl_ctx.compile(bufStructs, "md5.cl", "pbkdf2.cl") elif rtype == "sha1": - if saltlen < 32 and dklen < 32: + use_fast_kernel = ( + saltlen < 32 and dklen < 32 and max_password_bytes <= 32 + ) + if use_fast_kernel: dklen = 32 - self.max_out_bytes = bufStructs.specifySHA1(32, saltlen, dklen) + self.max_out_bytes = bufStructs.specifySHA1( + 32, + saltlen, + dklen, + max_password_bytes=32, + ) prg = self.opencl_ctx.compile(bufStructs, "pbkdf2_sha1_32.cl", None) else: - self.max_out_bytes = bufStructs.specifySHA1(128, saltlen, dklen) + max_in_bytes = max(128, max_password_bytes) + self.max_out_bytes = bufStructs.specifySHA1( + max_in_bytes, + saltlen, + dklen, + max_password_bytes=max_password_bytes, + ) prg = self.opencl_ctx.compile(bufStructs, "sha1.cl", "pbkdf2.cl") elif rtype == "sha256": - if saltlen <= 64 and dklen <= 64: + use_fast_kernel = ( + saltlen <= 64 and dklen <= 64 and max_password_bytes <= 32 + ) + if use_fast_kernel: dklen = 64 - self.max_out_bytes = bufStructs.specifySHA2(256, 128, saltlen, dklen) - if saltlen <= 64 and dklen <= 64: + max_in_bytes = max(128, max_password_bytes) + self.max_out_bytes = bufStructs.specifySHA2( + 256, + max_in_bytes, + saltlen, + dklen, + max_password_bytes=max_password_bytes, + ) + if use_fast_kernel: prg = self.opencl_ctx.compile(bufStructs, "pbkdf2_sha256_32.cl", None) else: prg = self.opencl_ctx.compile(bufStructs, "sha256.cl", "pbkdf2.cl") elif rtype == "sha512": - self.max_out_bytes = bufStructs.specifySHA2(512, 256, saltlen, dklen) + max_in_bytes = max(256, max_password_bytes) + self.max_out_bytes = bufStructs.specifySHA2( + 512, + max_in_bytes, + saltlen, + dklen, + max_password_bytes=max_password_bytes, + ) prg = self.opencl_ctx.compile(bufStructs, "sha512.cl", "pbkdf2.cl") else: assert "Error on hash type, unknown !!!" @@ -985,7 +1041,7 @@ def cl_pbkdf2_saltlist_init(self, type, pwdlen, dklen): bufStructs = buffer_structs() if type == "md5": self.max_out_bytes = bufStructs.specifyMD5( - max_in_bytes=128, + max_in_bytes=max(128, pwdlen), max_salt_bytes=128, dklen=dklen, max_password_bytes=pwdlen, @@ -994,7 +1050,7 @@ def cl_pbkdf2_saltlist_init(self, type, pwdlen, dklen): prg = self.opencl_ctx.compile(bufStructs, "md5.cl", "pbkdf2.cl") elif type == "sha1": self.max_out_bytes = bufStructs.specifySHA1( - max_in_bytes=128, + max_in_bytes=max(128, pwdlen), max_salt_bytes=128, dklen=dklen, max_password_bytes=pwdlen, @@ -1004,7 +1060,7 @@ def cl_pbkdf2_saltlist_init(self, type, pwdlen, dklen): elif type == "sha256": self.max_out_bytes = bufStructs.specifySHA2( hashDigestSize_bits=256, - max_in_bytes=128, + max_in_bytes=max(128, pwdlen), max_salt_bytes=128, dklen=dklen, max_password_bytes=pwdlen, @@ -1013,7 +1069,7 @@ def cl_pbkdf2_saltlist_init(self, type, pwdlen, dklen): elif type == "sha512": self.max_out_bytes = bufStructs.specifySHA2( hashDigestSize_bits=512, - max_in_bytes=256, + max_in_bytes=max(256, pwdlen), max_salt_bytes=128, dklen=dklen, max_password_bytes=pwdlen, diff --git a/lib/opencl_brute/opencl_information.py b/lib/opencl_brute/opencl_information.py index 835ac877..ade41f29 100644 --- a/lib/opencl_brute/opencl_information.py +++ b/lib/opencl_brute/opencl_information.py @@ -10,17 +10,31 @@ Refactored out of 'opencl.py' ''' -import pyopencl as cl +try: + import pyopencl as cl +except ImportError as _pyopencl_import_error: # pragma: no cover - exercised when PyOpenCL missing + cl = None +else: + _pyopencl_import_error = None + + +def _require_pyopencl(): + if cl is None: + raise ImportError( + "pyopencl is required for querying OpenCL information but is not installed" + ) from _pyopencl_import_error class opencl_information: def __init__(self): pass def printplatforms(self): + _require_pyopencl() for i,platformNum in enumerate(cl.get_platforms()): print('Platform %d - Name %s, Vendor %s' %(i,platformNum.name,platformNum.vendor)) def printfullinfo(self): + _require_pyopencl() print('\n' + '=' * 60 + '\nOpenCL Platforms and Devices') for i,platformNum in enumerate(cl.get_platforms()): print('=' * 60)