HackTheBox Digital-Safety-Annex Challenge

#!/usr/bin/env python3
from multiprocessing import Pool, cpu_count
import sys, time, gmpy2
from pwn import *

def check_k_range(args):
    start, end, g, p, q, r = args
    for k in range(start, end):
        if k % 2000 == 0:
            print(f"Trying k = {k}")
        if pow(g, k, p) % q == r:
            return k
    return None

def find_k(g, p, q, r):
    cores = cpu_count()
    start = 65500
    end = 10**6
    chunk_size = (end - start) // cores
    ranges = []
    for i in range(cores):
        chunk_start = start + (i * chunk_size)
        chunk_end = chunk_start + chunk_size if i < cores - 1 else end
        ranges.append((chunk_start, chunk_end, g, p, q, r))
    with Pool(cores) as pool:
        results = pool.map(check_k_range, ranges)
    for result in results:
        if result is not None:
            return result
    return None

def exploit(target):
    ip, port = target.split(':')
    port = int(port)
    context.log_level = 'debug'
    io = remote(ip, port)
    io.recvuntil(b'>')
    io.sendline(b'4')
    buf = io.recvuntil(b'[+] Test user log (y/n) : ')
    p = int(buf.decode().split('\n')[6].split(' = ')[1])
    q = int(buf.decode().split('\n')[7].split(' = ')[1])
    g = int(buf.decode().split('\n')[8].split(' = ')[1])
    io.sendline(b'y')
    io.recvuntil(b'Enter your password : ')
    io.sendline(b'5up3r_53cur3_P45sw0r6')
    buf = io.recvuntil(b'>')
    r = int(buf.decode().split('\n')[1].split('((')[6].split(',')[0])
    s = int(buf.decode().split('\n')[1].split('((')[6].split(',')[1][1:-1])
    h = int(buf.decode().split('\n')[1].split('((')[6].split(',')[2].split(')]')[0].split("'")[1], 16)
    print(f"Starting k search using {cpu_count()} cores...")
    kx = find_k(g, p, q, r)
    if not kx:
        print("Failed to find k value")
        return
    print(f"Found k = {kx}")
    io.sendline(b'3')
    io.recvuntil(b'Please enter the username who stored the message : ')
    io.sendline(b'ElGamalSux')
    io.recvuntil(b'Please enter the message\'s request id: ')
    io.sendline(b'3')
    io.recvuntil(b'Please enter the message\'s nonce value : ')
    io.sendline(str(kx).encode())
    io.recvuntil(b'[+] Please enter the private key: ')
    pri = ((kx * s - h) * gmpy2.invert(r, q)) % q
    io.sendline(str(pri).encode())
    response = io.recvall().decode()
    print(response)

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(f"Usage: {sys.argv[0]} <ip:port>")
        sys.exit(1)
    target = sys.argv[1]
    exploit(target)