import csv
import random

random.seed(42)

SQLI_TEMPLATES = [
    "' OR '1'='1",
    "' OR 1=1 --",
    "admin'--",
    "' UNION SELECT username, password FROM users--",
    "1; DROP TABLE users;--",
    "' AND SLEEP(5)--",
    "\" OR \"\"=\"",
    "1' AND '1'='1' UNION SELECT NULL,NULL--",
    "'; EXEC xp_cmdshell('dir')--",
    "' OR 'a'='a",
    "1 OR 1=1",
    "' OR EXISTS(SELECT * FROM users)--",
    "%27%20OR%20%271%27%3D%271",
    "'/**/OR/**/1=1--",
]

BENIGN_TEMPLATES = [
    "john.doe@example.com",
    "Jane O'Brien",
    "New York, NY",
    "Sunset Boulevard Apt 4B",
    "search term: winter jackets",
    "user123",
    "2024-05-19",
    "5 stars, would buy again",
    "quantity: 3",
    "order #{n}",
    "product review for {n}",
    "phone: 555-01{n}",
    "comment: great service, thanks!",
    "coupon code SAVE{n}",
]


def fill(template):
    return template.format(n=random.randint(10, 99))


def generate(n_per_class=250):
    rows = []
    for _ in range(n_per_class):
        rows.append((random.choice(SQLI_TEMPLATES), 1))
        rows.append((fill(random.choice(BENIGN_TEMPLATES)), 0))
    random.shuffle(rows)
    return rows


if __name__ == "__main__":
    rows = generate(n_per_class=250)
    with open("data.csv", "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["text", "label"])
        writer.writerows(rows)
    print(f"Wrote {len(rows)} rows to data.csv")