CTF中的Speedup

import hashlib
import itertools
from multiprocessing import Pool, cpu_count, current_process, set_start_method
from tqdm import tqdm

set_start_method("fork")


def sm3_hash(msg):
    hash_obj = hashlib.new("sm3")
    hash_obj.update(msg)
    return hash_obj.hexdigest()


def attack(bound):
    low, up = bound

    search_space = itertools.product(range(low, up), *[range(256)] * 3)

    for i, j, k, l in search_space:
        tmp = block[i] + block[j] + block[k] + block[l]
        msg = m + tmp.encode()

        if sm3_hash(msg) == h:
            print(f"\n!!! Flag Found by {current_process().name}: {msg}")
            return msg

    return None


if __name__ == "__main__":
    table = ["a", "b", "c", "d"]
    m = b"adcddbbadcacabad"
    h = "aab05fca300811223b3b957bfe33130770fb7a6b55b030a5809c559344f66f79"
    block = ["".join(i) for i in itertools.product(table, repeat=4)]
    num_tasks = cpu_count() * 4
    total_range = 256
    chunk_size = (
        total_range // num_tasks
        if total_range % num_tasks == 0
        else total_range // num_tasks + 1
    )
    bounds = []
    for i in range(num_tasks):
        low = i * chunk_size
        up = min((i + 1) * chunk_size, total_range)
        if low < up:
            bounds.append((low, up))

    print(f"Starting attack with 10 processes, divided into {len(bounds)} tasks...")

    with Pool(10) as pool:
        results_iterator = pool.imap_unordered(attack, bounds)

        for result in tqdm(results_iterator, total=len(bounds), desc="Cracking Hash"):
            if result:
                pool.terminate()
                pool.join()
                break

    print("Attack finished.")

import numpy as np
from numba import njit, objmode


@njit
def _h_jit(b: np.ndarray) -> np.uint64:
    h1 = np.uint64(0x1234567890AB)
    h2 = np.uint64(0xFEDCBA098765)

    for i in range(len(b)):
        byte = b[i]
        shift = np.uint64((i % 6) * 6)

        if i % 2 == 0:
            h1 ^= np.uint64(byte) << shift
            h1 = (h1 * np.uint64(0x100000001B3)) & np.uint64(0xFFFFFFFFFFFF)
        else:
            h2 ^= np.uint64(byte) << shift
            h2 = (h2 * np.uint64(0xC6A4A7935BD1)) & np.uint64(0xFFFFFFFFFFFF)

    with objmode(result="uint64"):
        py_h1 = int(h1)
        py_h2 = int(h2)
        py_res = py_h1 ^ ((py_h2 << 24) | (py_h2 >> 24))
        py_res = (py_res ^ (py_res >> 25)) * 0xFF51AFD7ED55
        py_res &= 0xFFFFFFFFFFFFFFFF
        result = np.uint64(py_res)
    result = (result ^ (result >> np.uint64(25))) * np.uint64(0xC4CEB9FE1A85)
    result ^= result >> np.uint64(25)

    return result


def h(data: str) -> bytes:
    b_np = np.frombuffer(data.encode("utf-8"), dtype=np.uint8)
    result_int = _h_jit(b_np)
    return result_int.to_bytes(8, "big")[:6]

package main

import (
	"fmt"
	"runtime"
	"sync"
	"sync/atomic"

	"github.com/schollz/progressbar/v3"
	"github.com/tjfoc/gmsm/sm3"
)

const (
	targetHash = "aab05fca300811223b3b957bfe33130770fb7a6b55b030a5809c559344f66f79"
	prefix     = "adcddbbadcacabad"
)

var charTable = []byte("abcd")

func worker(wg *sync.WaitGroup, jobs <-chan uint64, found *int32, result *string, resultLock *sync.Mutex, bar *progressbar.ProgressBar) {
	defer wg.Done()

	prefixBytes := []byte(prefix)

	for j := range jobs {
		if atomic.LoadInt32(found) == 1 {
			bar.Add(1)
			return
		}

		var i uint64
		for i = 0; i < (1 << 24); i++ {
			fullIndex := (j << 24) + i

			suffix := make([]byte, 16)
			for k := 0; k < 16; k++ {
				shift := 30 - (k * 2)
				charIndex := (fullIndex >> shift) & 3
				suffix[k] = charTable[charIndex]
			}

			data := append(prefixBytes, suffix...)

			hash := sm3.New()
			hash.Write(data)
			sum := hash.Sum(nil)
			sumHex := fmt.Sprintf("%x", sum)

			if sumHex == targetHash {
				if atomic.CompareAndSwapInt32(found, 0, 1) {
					resultLock.Lock()
					*result = string(data)
					resultLock.Unlock()
				}
				bar.Add(1)
				return
			}
		}
		bar.Add(1)
	}
}

func main() {
	numWorkers := runtime.NumCPU()
	fmt.Printf("Starting attack with %d workers...\n", numWorkers)

	var wg sync.WaitGroup
	var found int32
	var result string
	var resultLock sync.Mutex

	jobs := make(chan uint64, 256)

	bar := progressbar.NewOptions(256,
		progressbar.OptionSetDescription("Cracking Hash"),
		progressbar.OptionSetPredictTime(true),
		progressbar.OptionShowCount(),
		progressbar.OptionSetTheme(progressbar.Theme{
			Saucer:        "=",
			SaucerHead:    ">",
			SaucerPadding: " ",
			BarStart:      "[",
			BarEnd:        "]",
		}),
	)

	wg.Add(numWorkers)
	for w := 0; w < numWorkers; w++ {
		go worker(&wg, jobs, &found, &result, &resultLock, bar)
	}

	for j := uint64(0); j < 256; j++ {
		jobs <- j
	}
	close(jobs)

	wg.Wait()
	bar.Finish()

	if result != "" {
		fmt.Printf("\n\n!!! Flag Found: %s\n", result)
	} else {
		fmt.Println("\n\nAttack finished. Message not found.")
	}
}