https://app.hackthebox.com/challenges/710

Description

After many perils and challenges, you have finally entered vault 79. Going through the remnants of a place people once dwelled by people, deep inside the vault you lay your eyes upon an unsealed bunker. You move it’s entrance, and inside you lay your eyes on the long-sought gold. As you celebrate and rejoice, the door behind you closes, trapping you inside. A screen appears, telling you that you need to give big amounts of gold based on numbers in the screen, or you will be stuck inside forever. After solving so many problems however, and so close to success, you aren’t prepared to give away the gold just to escape. You start studying what the screen says, and start to devise a way to keep as much gold as you can. Can you give away the least possible gold, and outsmart the vault?

Exploitation

#!/usr/bin/env python3
from pwn import *

def get_process():
    try:
        host, port = sys.argv[1].split(':')
        return remote(host, int(port))
    except IndexError:
        print(f'Usage: python {sys.argv[0]} <ip:port>')
        exit(1)

def get_values(test_n):
    global a
    io.recvuntil(f'Test {test_n + 1}/100\n'.encode())
    n, k = list(map(int, io.recvline().rstrip().decode().split()))
    a = [0] + list(map(int, io.recvline().rstrip().decode().split()))
    return n, k

def divide_and_conquer(i, l, r, optl, optr):
    if l > r:
        return
    mid = (l + r) // 2
    best_k = -1
    dp[i][mid] = float('inf')
    global cl, cr, cost, count
    for k in range(optl, min(mid, optr) + 1):
        while cr < mid:
            cr += 1
            add(cr)
        while cr > mid:
            remove(cr)
            cr -= 1
        while cl < k + 1:
            remove(cl)
            cl += 1
        while cl > k + 1:
            cl -= 1
            add(cl)
        tmp = dp[i - 1][k] + cost
        if tmp < dp[i][mid]:
            dp[i][mid] = tmp
            best_k = k
    divide_and_conquer(i, l, mid - 1, optl, best_k)
    divide_and_conquer(i, mid + 1, r, best_k, optr)

def find_minimum_cost(n, k):
    global dp, cl, cr, cost, count
    maxA = max(a)
    dp = [[float('inf')] * (n + 1) for _ in range(k + 1)]
    dp[0][0] = 0
    for i in range(1, k + 1):
        cl, cr = 1, 0
        cost = 0
        count = [0] * (maxA + 1)
        divide_and_conquer(i, i, n, i - 1, n - 1)
    return dp[k][n]

def add(pos):
    global cost
    x = a[pos]
    cost += count[x]
    count[x] += 1

def remove(pos):
    global cost
    x = a[pos]
    count[x] -= 1
    cost -= count[x]

def send_solution(min_cost):
    io.sendline(f'{int(min_cost)}'.encode())

def get_flag():
    io.recvuntil(b'HTB{')
    return b'HTB{' + io.recvline().rstrip()

def pwn():
    for t in range(100):
        print('Test', t + 1)
        n, k = get_values(t)
        min_cost = find_minimum_cost(n, k)
        send_solution(min_cost)
    flag = get_flag()
    print(flag)

if __name__ == '__main__':
    a = []
    dp = []
    cost = 0
    count = []
    cl = 0
    cr = 0
    io = get_process()
    pwn()

Summary

Nothing Without A Cost: reduce the custom rules to a scriptable check and use the smallest reliable path to the flag.