parameterize.py

#!/usr/bin/python3

import argparse, server
from itertools import product

description = '''Parameterize a UUID:
\tpython3 parameterize.py -c "abcdef0123456789-" -l 36
Parameterize a random sentence:
\tpython3 parameterize.py -s "hello world!"
Parameterize a random sentence to exfiltrate from id-attribute of an img-element:
\tpython3 parameterize.py -s "hello world!" -a "id" -e "img"'''

categories = {
    "numbers": "0123456789",
    "alphabet": "abcdefghijklmnopqrstuvwxyz",
    "alphabet_upper": "ABCDEFGHIJKLMNOPQRSTUVWXYZ",
    "german": "öäüß",
    "german_upper": "ÖÄÜ"
}

# Initialize argument parsing
parser = argparse.ArgumentParser(
    epilog=description,
    formatter_class=argparse.RawDescriptionHelpFormatter,
    description="CLI tool to parameterize the exfiltration server")

general_group = parser.add_argument_group("INPUT")
parser.add_argument("-s", "--sample", type = str, help="A sample secret")
parser.add_argument("-c", "--charset", type = str, help="All possible characters that can occur (overwrites -a parameter)")

scan_group = parser.add_argument_group('PROPERTIES')
parser.add_argument("-l", "--length", type = int, default=0, help="Fixed length of the secret")
parser.add_argument("-a", "--attribute", type = str, default="src", help="Attribute name to exfiltrate (e.g. src, href)")
parser.add_argument("-e", "--element", type = str, default="img", help="Element, class or id to exfiltrate (e.g. img, div, #id, .class)")
parser.add_argument("--allchars", action='store_true', help="If a sample is given, all characters of the same categories will be added (if sample contains '1', all other numbers between 0 and 9 will be included)")
args = parser.parse_args()

# Assertions
assert isinstance(args.sample, str) or (isinstance(args.length, int) and args.length != 0), "No sample or length was supplied"
if args.length != 0 and args.sample: assert args.length == len(args.sample), "Length parameter mismatches sample length"
secret_length = args.length == 0 and len(args.sample) or args.length
assert secret_length >= 3, "Length must be >= 3"

charset, occuring_chars = None, None
if args.sample: occuring_chars = set([*args.sample])
if args.charset: charset = set([*args.charset])
assert charset or occuring_chars, "No sample or charset was provided"

# Check if sample contains a 3-char combination twice or more
# Example: "abcXYZdefghXYZijkl" contains "XYZ" twice
if args.sample:
    for i in range(len(args.sample) - 2):
        assert args.sample[i:i+3] not in args.sample[i+3:], "Sample contains a 3-char combination twice or more"

if charset and occuring_chars:
    assert len(occuring_chars - charset) == 0, "Sample contains characters not inluded inside the charset"

if not charset:
    if args.allchars:
        found_categories = set()
        # This works, but is not elegant
        for char in occuring_chars.copy():
            for category, value in categories.items():
                if char in value and not category in found_categories:
                    occuring_chars |= set([*value])
                    found_categories.add(category)
    charset = occuring_chars

# Calculate all 3-char combinations
char_combinations = set()
for i in product(list(charset), repeat=3):
    char_combinations.add(''.join(map(str, i)))
combination_count = len(char_combinations)

if combination_count > 20_000:
    print("WARNING: Selector count is higher than 20.000. This will take some time...")

# Calculate size of CSS file for the client to download
total_size = len(server.generate_payload(''.join(charset), element=args.element, attribute=args.attribute))

# Convert charset back to list and sort
charset = sorted(list(charset), key = lambda s: sum(map(ord, s)))

# Edgecase: if first character is a '-', put it at the end
if charset[0] == '-':
    charset = charset[1:] + ['-']

print(f"Charset: '{''.join(charset)}'")
print("Secret length: ", secret_length)
print("Number of selectors: ", combination_count)
print(f"Filesize of style.css: {total_size} Bytes")

# Calculate number of requests the client has to make
num_requests = secret_length - 2
print("Number of client requests required: ", num_requests)