llama2_mutator.py

# PPO process that receive seeds from fuzzer then send back the mutated seeds to fuzzer
from dataclasses import dataclass, field
from typing import Optional
import os
import torch
import tyro
from accelerate import Accelerator
from peft import LoraConfig
from tqdm import tqdm
from transformers import AutoTokenizer, BitsAndBytesConfig

from trl import (
    AutoModelForCausalLMWithValueHead,
    PPOConfig,
    set_seed,
)
import threading
import re
import sysv_ipc
import struct
import random

tqdm.pandas()

TYPE_SEED = 1
TYPE_TEXT_SEED = 2
TYPE_REWARD = 3
TYPE_REQUEST = 4

access_token = "YOUR ACCESS TOKEN"
cur_path = os.path.dirname(os.path.realpath(__file__))
output_dir = os.path.join(cur_path, "ppo_checkpoint")
message_queue = []
seed_id_map = {}
id_rwd_map = {}
seeds_from_fuzzer = set()
uid = 1
shared_resource_lock = threading.Lock()


@dataclass
class ScriptArguments:
    """
    Setup experiment config
    """

    fuzzing_target: Optional[str] = field(default="libpng")
    fuzzing_object: Optional[str] = field(default="")
    if_mixed_model: Optional[bool] = field(default=True)
    if_text: Optional[bool] = field(default=False)
    temperature: Optional[float] = field(default=1.25)
    peft_config: Optional[LoraConfig] = field(
        default_factory=lambda: LoraConfig(
            r=64,
            lora_alpha=16,
            lora_dropout=0.1,
            bias="none",
            target_modules=[
                "q_proj",
                "down_proj",
                "gate_proj",
                "o_proj",
                "k_proj",
                "v_proj",
                "up_proj",
            ],
            task_type="CAUSAL_LM",
        ),
    )
    trust_remote_code: bool = field(
        default=True, metadata={"help": "Enable `trust_remote_code`"}
    )


args = tyro.cli(ScriptArguments)


def mq_thread():
    """
    Thread to receive request from fuzzer, and send generated seed to fuzzer
    """
    global message_queue, seed_id_map, seeds_from_fuzzer
    try:
        mq = sysv_ipc.MessageQueue(1234, sysv_ipc.IPC_CREAT)
    except sysv_ipc.ExistentialError:
        print(f"Message queue with key {1234} already exists.")
        return
    while True:
        # only receive request msg
        try:
            msg, mtype = mq.receive(type=TYPE_REQUEST)
            if msg != b"":
                if len(seeds_from_fuzzer) > 30:
                    seeds_from_fuzzer.clear()
                seeds_from_fuzzer.add(msg.decode(errors="ignore")[4:])
            if message_queue != []:
                # send uid + seed
                seed = message_queue.pop(0)
                send_msg = struct.pack("I", seed_id_map[seed])
                if len(seed) > (2045 - len(send_msg)):
                    seed = seed[: (2045 - len(send_msg))]
                send_msg = send_msg + seed.encode("utf-8")
                if len(send_msg) >= 2045:
                    print("::oversize")
                    continue
                if not args.if_text:
                    mq.send(
                        send_msg,
                        True,
                        type=TYPE_SEED,
                    )
                else:
                    mq.send(
                        send_msg,
                        True,
                        type=TYPE_TEXT_SEED,
                    )
        except RuntimeError as e:
            print(e)


def hex_string_to_hex(hex_string, fuzzing_target, if_text):
    """
    Formatting generated hex string.

    Returns:
        String of hex.
    """
    if len(hex_string.split("### Output:")) >= 2:
        hex_string = hex_string.split("### Output:")[1]
    else:
        hex_string = hex_string.replace(
            f"### Input: ```Based on below hex {fuzzing_target} seed, mutate a new {fuzzing_target} seed. Make sure the example is complete and valid.",
            " ",
        )
    if not if_text:
        hex_string = re.sub(r"[^a-zA-Z0-9\s]", " ", hex_string)
        hex_values = hex_string.replace("0x", " ")
        # Split the string into sections
        sections = hex_values.split()
        # Iterate through the sections and add leading zeros if needed
        result = []
        for section in sections:
            if len(section) == 1:
                section = "0" + section
                result.append(section)
            elif len(section) == 2:
                result.append(section)
        result = "".join(result)
    else:
        result = hex_string
    if len(result) > 2040:  # limite seed size to 2048
        result = result[:2040]
    return result


def main():
    """
    Main function to run PPO loop
    """
    model_name = f"llama-2-7b-structured-{args.fuzzing_target}-hex-mutator"
    if args.if_mixed_model:
        model_name = f"llama-2-7b-structured-{args.fuzzing_target}-mix-hex-mutator"
    if args.fuzzing_object != "":
        model_name = f"llama-2-7b-structured-{args.fuzzing_target}-{args.fuzzing_object}-mix-hex-mutator"
        args.fuzzing_target = args.fuzzing_target + " " + args.fuzzing_object
    # Init the tokenizer and dataset
    tokenizer = AutoTokenizer.from_pretrained(
        os.path.join(cur_path, model_name),
        use_fast=True,
        token=access_token,
    )
    # Some tokenizers like GPT-2's don't have a padding token by default, so we set one here.
    # tokenizer.pad_token_id = tokenizer.eos_token_id
    tokenizer.pad_token = tokenizer.bos_token
    # tokenizer.padding_side = "left"
    # We retrieve the dataloader by calling the `build_dataset` function.

    # set seed before initializing value head for deterministic eval
    set_seed(0)

    # Build the model.
    peft_config = args.peft_config
    # Copy the model to each device
    current_device = Accelerator().local_process_index
    device_map = {"": current_device}

    model = AutoModelForCausalLMWithValueHead.from_pretrained(
        os.path.join(cur_path, model_name),
        trust_remote_code=args.trust_remote_code,
        device_map=device_map,
        peft_config=peft_config,
        token=access_token,
        torch_dtype=torch.bfloat16,
        quantization_config=BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        ),
        # use_flash_attention_2=True, Unable to use this feature in current GPU
    )
    # Whether or not the model should use the past last key/values attentions (if applicable to the model) to speed up decoding.
    model.config.use_cache = False
    model.config.pretraining_tp = 1

    # flash attention 1
    torch.backends.cuda.sdp_kernel(
        enable_flash=True, enable_math=False, enable_mem_efficient=False
    )
    example = {
        "libpng": "0x890x50,0x4e0x47,0xd0xa,0x1a0xa,0x00x0,0x00xd,0x490x48,0x440x52,0x00x0,0x00x1,0x00x0,0x00x1,0x20x3,0x00x0,0x10x3e,0xb30xd8,0x210x0,0x00x0,0x60x50,0x4c0x54,0x450xee,0xff0x22,0x220x66,0xff0x6c,0x20xd2,0x260x0,0x00x0,0x290x49,0x440x41,0x540x8,0xd70x63,0x600x84,0x820xf,0x500x0,0xd0x49,0x480x44,0x580x4e,0x900xe0,0x20xd5,0x70x89,0x500x4e,0x470xd,0xa0x1a,0xa0x0,0x00x0,0xd0x49,0x480x44,0x520x0,0x00x0,0x200x0,0x00x0,0x20x1,0x30x0,0x00x1,0x3e0xb3,0xd80x21,0x00x0,0x00x6",
        "libtiff": "0x490x49,0x2a0x0,0x620x0,0x00x0,0x740xa5,0xc20xae,0xce0xa,0xbf0x9,0x130x42,0x580x2d,0x2b0xbf,0xee0x20,0xb70x1b,0x430x5f,0xff0x23,0x840xb0,0xb0xed,0x950xb1,0x5f0xf5,0x9b0x70,0x530x5b,0x200x62,0x90xca,0x530x7a,0x260x39,0x7e0xe4,0xfa0xf,0xf0xf,0xf0xf,0xa80x6e,0x6e0x2f,0xc40x49,0x300xc8,0x3d0x55,0x7e0x4e,0x880x6e,0x500x1,0x9f0x64,0x690x6b,0x5d0xa8,0xd20xe9,0x480x88,0xbf0xfb,0x2f0xf,0xae0x0,0x00x0,0x00x30,0x6b0x1b,0x810x1,0x120x0,0x00x1,0x30x0,0x10x0,0x00x0,0x30x0,0x10x1,0x10x1,0x30x0,0x10x0,0x00x0,0x1d0x0,0x00x0,0x20x1,0x30x0,0x10x0,0x00x0,0x100x0,0x00x0,0x30x1,0x30x0,0x10x0,0x00x0,0x80x0,0x00x0,0x150x1,0x30x0,0x10x0,0x00x0,0x60x1,0x10x0,0x3d0x1,0x30x0,0x10x0,0x00x0,0x20x0,0x00x0,0xd0x1,0x20x0,0xf0x0,0x00x0,0x400x1,0x00x0,0x110x1,0x40x0,0x10x0,0x00x0,0x170x1,0x00x0,0xd0x0,0x1c0x0,0x10x0,0x00x0,0x10x0,0x00x1,0xa0x0,0x30x0,0x10x0,0x00x0,0x20x1,0x10x0,0x160x1,0x30x0,0x10x0,0x00x0,0x10x0,0xec0x0,0xa0x0,0x1c0x1,0x00x0,0x00x0,0x00x1,0x00x1,0x160x1,0x1c0x1,0x00x1,0x10x1,0x90x1,0x00x1,0x170x1,0x40x0,0x00x1,0x10x0,0x800x0,0x00x0,0x1a0x1,0x80x0,0x10x0,0x00x0,0x480x1,0x00x1,0x1b0x1,0x50x0,0x10x0,0x00x0,0x5a0x0,0x00x0,0x1c0x1,0x30x0,0x10x0,0x00x0,0x10x0,0x00x0,0x280x1,0x30x0,0x10x0,0x00x0,0x10x0,0x10x0,0x240x0,0x1c0x1,0x20x0,0x10x1,0x10x0,0x10x0,0x00x0,0x10x1,0x620xc2,0xc20xc2",
        "poppler": "0x250x50,0x440x46,0x2d0x31,0x2e0x34,0xa0x31,0x200x30,0x200x6f,0x620x6a,0x200xa,0x3c0x3c,0xa0x2f,0x500x61,0x670x65,0x730x20,0x320x20,0x300x20,0x520xa,0x320x20,0x300x20,0x6f0x62,0x6a0x20,0xa0x3c,0x3c0xa,0x2f0x52,0x650x73,0x6f0x75,0x720x63,0x650x73,0x200xa,0x3e0x3e,0xa0x2f,0x430x6f,0x6e0x74,0x650x6e,0x740x73,0x200x38,0x310x31,0x200x30,0x200x52,0xa0x38,0x310x31,0x200x30,0x200x6f,0x620x6a,0x200xa,0x3c0x3c,0xa0x2f,0x4c0x65,0x6e0x67,0x740x68,0x200x31,0x370x38,0x360x33,0xa0x3e,0x3e0xa,0x730x74,0x720x65,0x610x6d,0xa0x42,0x490xa,0x2f0x57,0x200x36,0x320xa,0x2f0x48,0x200x36,0x320xa,0x2f0x44,0x5b0x31,0xa0x30,0x5d0xa,0x2f0x46,0x2f0x43,0x430x46,0x8a0x2f,0x440x50,0x3c0x3c,0x2f0x4b,0x200x2d,0x310xa,0x740x72,0x610x69,0x6c0x65,0x720xa,0x3c0x3c,0xa0x2f,0x520x6f,0x6f0x74,0x200x31,0x200x30,0x200x52,0xa0x3e,0x3e0xa",
        "openssl asn1": "0x300x80,0x300x80,0x60x9,0x2a0x86,0x480x86,0xf70xd,0x10x7,0x20xa0,0x800x30,0x800x2,0x10x80,0x310x0,0x300x80,0x60x9,0x2a0x86,0x480x86,0xf70xd,0x10x7,0x20xa0,0x800x30,0x800x2,0x10x80,0x310x0,0x300x80,0x60x9,0x2a0x86,0x480x86,0xf70xd,0x10x7,0x20xa0,0x800x30,0x800x2,0x10x80,0x310x0,0x300x80,0x60x9,0x2a0x86,0x480x86,0xf70xd,0xcc0xa9,0x720x5f,0x920x9f,0x290x96,0xd50x34,0xe70x7,0x90x99,0x780xf5,0x3d0xe0,0x290x24,0x130x8b,0x870x7e,0xfa0xf0,0x810xfb,0x70xf6,0xce0x4d,0x5b0xe8,0x940xe4,0xc90xd8,0xe90x6f,0x420xa8,0x9c0x58,0x5c0x8d,0x430xad,0x3e0x53,0x2d0x68,0xc30xd1,0x640xbb,0x720x27,0xbc0x2a,0x90x82,0x4f0xce,0x3e0x81,0x6c0x3e,0xe00x31,0xe90x9c,0xef0x75,0xfa0x4c,0x310x0,0x300x80,0x60x9,0x2a0x86,0x480x86,0xf70xd,0x10x7,0x20xa0,0x800x30,0x800x2,0x10x80,0x310x0,0x300x80,0x60x9,0x2a0x86,0x480x86,0xf70xd,0x10x7,0x20xa0,0x800x30,0x800x2,0x10x80,0x310x0,0x300x80,0x60x9,0x2a0x86,0x480x86,0xf70xd,0x10x7,0x20xa0,0x800x30,0x800x2,0x10x80,0x310x0,0x300x80,0x60x9,0x2a0x86,0x480x86,0xf70xd,0x10x7,0x20xa0,0x800x30,0x800x2,0x10x80,0x310x0,0x300x80,0x60x9,0x2a0x86,0x480x86,0xf70xd,0x10x7,0x20xa0,0x800x30,0x800x2,0x10x80,0x310x0,0x300x80,0x60x9,0x2a0x86,0x480x86,0xf70xd,0x10x7,0x20xa0,0x800x30,0x800x2,0x10x80,0x310x0,0x300x80,0x60x2,0x300x30,0x00x0,0xa10x80,0x300x80,0x300x80,0x300x80,0x350x1,0x300x0,0x0",
        "openssl asn1parse": "0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30,0x120x1,0x300x12,0x10x30",
        "openssl bignum": "0x30x4,0xff0x0,0x10xf,0xfe0xff,0x50x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x50x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x3a,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x1,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0xa40xb,0xf70x7,0xa50x26,0x700x4b,0x890xc7,0x870xf5,0xfb0x71,0xd00x73,0x5a0x99,0x500x3c,0x8e0xd2,0x590x5e,0x00x6a,0x520xf0,0xcd0x40,0xee0x1b,0xf10xa6,0x200x85,0xe50x87,0x150xc2,0xe00xb0,0x4f0xdd,0xf60xbc,0x900x70,0xfc",
        "openssl x509": "0x300x80,0x300x80,0x20x2,0x300x30,0x300x80,0x60x2,0x2b0x3,0x00x0,0x300x0,0x300x80,0x180x1,0x3a0x18,0x10x30,0x00x0,0x300x0,0x300x82,0x10x2a,0x300x81,0xeb0x6,0x70x2a,0x860x48,0xce0x3d,0x20x1,0x300x81,0xdf0x2,0x10x1,0x300x28,0x60x7,0x2a0x86,0x480xce,0x3d0x1,0x10x2,0x1d0x0,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x1,0x300x53,0x40x1c,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xfe,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xfe,0x40x1c,0xb40x5,0xa0x85,0xc0x4,0xb30xab,0xf50x41,0x320x56,0x500x44,0xb00xb7,0xd70xbf,0xd80xba,0x270xb,0x390x43,0x230x55,0xff0xb4,0x30x15,0x00xbd,0x710x34,0x470x99,0xd50xc7,0xfc0xdc,0x450xb5,0x9f0xa3,0xb90xab,0x8f0x6a,0x940x8b,0x140x4,0x390x4,0xb70xe,0xc0xbd,0x6b0xb4,0xbf0x7f,0x320x13,0x900xb9,0x4a0x3,0xc10xd3,0x560xc2,0x110x22,0x340x32,0x800xd6,0x110x5c,0x1d0x21,0xbd0x37,0x630x88,0xb50xf7,0x230xfb,0x4c0x22,0xdf0xe6,0xcd0x43,0x750xa0,0x5a0x7,0x470x64,0x440xd5,0x810x99,0x850x0,0x7e0x34,0x20x1d,0x00xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xf0xfc,0x4f0x0,0x00x0,0x00x0,0x00x0,0x10x2,0x10x1,0x30x3a,0x00x4,0x6a0xf2,0x640xdc,0xe30x33,0x3e0xe5,0x300x81,0xdf0x0,0x10x30,0x300x28,0x60x7,0x2a0x86,0x480xce,0x3d0x1,0x10x2,0x1d0x0,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0xff0xff,0x00x0,0x00x0,0x00x0,0x00x0,0x00x1,0xff0x2",
        "openssl client": "0x160x3,0x30x2,0x130x2,0x00x0,0x510x3,0x00x0,0x00x0,0x00x0,0x20x2a,0x760x72,0x230xc9,0x9a0x7,0x350x7,0xc0x9,0x720x19,0x60xc5,0x8f0x27,0xea0x6e,0x100x41,0xcb0xa5,0xaf0xf6,0xdc0x20,0xbe0x32,0xb30x4a,0xd40x45,0xf30x94,0xe10x63,0xcf0xec,0xa50x38,0x610x45,0x5b0xd2,0x6d0x57,0x420xca,0xe60x98,0x280xc9,0xc0x9b,0x840xe5,0x10x0,0x00x1,0x00x0,0x90xff,0x10x0,0x10x0,0x00x0,0x00x0,0xb0x0,0x10xb6,0x00x1,0xb30x0,0x10xb0,0x300x82,0x10xac,0x300x82,0x10x15,0xa00x3,0x20x1,0x20x2,0x10x2,0x300xd,0x60x9,0x2a0x86,0x480x86,0xf70xd,0x10x1,0xb0x5,0x00x30,0xe0x31,0xc0x30,0xa0x6,0x30x55,0x40x3,0x130x3,0x720x73,0x610x30,0x1e0x17,0xd0x31,0x360x30,0x350x31,0x390x31,0x370x30,0x320x34,0x310x5a,0x170xd,0x320x36,0x300x39,0x310x37,0x310x35,0x300x32,0x340x31,0x5a0x30,0xe0x31,0xc0x30,0xa0x6,0x30x55,0x40x3,0x130x3,0x720x73,0x610x30,0x810x9f,0x300xd,0x60x9,0x2a0x86,0x480x86,0xf70xd,0x10x3,0x10x30,0x00x3,0x810x8d,0x00x30,0x810x89,0x20x81,0x810x0,0x00x0,0x00x0,0x00x0,0x00x20,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x100x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x54,0x00x0,0x00x1,0x00x0,0x60x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x40,0x00x0,0x00x0,0x00x80,0x00x0,0x20x0,0x00x0,0x00x0,0x00x0,0x00xf1,0x00x0,0x00x0,0x00x0,0x70x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0x00x0,0xfd0x2,0x30xf7,0xff0x1,0xa30x1a,0x310x18,0x300x9,0x60x3,0x550x1d,0x130x0,0x300xb,0x60x3,0x550x1d,0xf0x4,0x40x3,0x20x5,0xa00x30,0xd0x6,0x90x2a,0x860x48,0x860xf7,0xe50xe5,0xe50xd,0x10x1,0xb0x5,0x00x3,0x810x81,0x00x88,0xf90x80,0x340xca,0xd40xf1,0x1b0xbe,0x810xa0,0x1c0x64,0x580x84,0xbb0x3c,0x320x3b,0xad0x58,0xc90xcd,0x340xf7,0x6d0x1e,0xee0x90,0x640x15,0x880xb5,0x440x71,0xdd0x95,0xd60x3b,0x70x4e,0x2a0x34,0x960x5c,0xa00x35,0xb20x37,0xb20xce,0xd60x3e,0x750x2f,0xe40xe9,0xf60x3,0xa30xb9,0x5c0xe2,0x470x45,0xfa0x1f,0x10x0,0x00x20,0x610x11,0xa20x89,0x650x25,0x620xff,0x250x8d,0x240xc9,0x2d0x50,0x520x3a,0xe0x4d,0x40x9e,0x60x3f,0xa10xfe,0xce0x71,0x720x24,0xf30xcf,0x250x62,0x9f0x5,0x170x9e,0xd40xe9,0xc10x7a,0xac0x17,0x710xba,0x840x74,0x740x7b,0xcf0x32,0xaf0xaa,0x9d0x94,0xd80xaf,0x8e0xe,0x00x0,0x00x14,0xbe0x0,0x160x0,0xaa0xc,0x00x3,0x730x0,0x10x10,0xe70x3,0xfe0x14,0x70x2c",
        "openssl server": "0x160x3,0x10x1,0x3b0x1,0x00x1,0x370x3,0x20xf4,0x400x36,0xa50x59,0x300x5,0x250x51,0xd70x3c,0x190x57,0x480x96,0xb0xdc,0xf00xda,0x9b0xe4,0x650xe0,0x660x27,0xa90xc5,0xc60xe8,0x840xe2,0x5f0x10,0xe60x74,0x220x12,0xab0x45,0xbc0x4d,0xdc0x8a,0x70x6d,0x150xee,0x1b0x77,0x00x20,0x00x0,0x00x0,0x00x0,0x10xd9,0xc00x16,0xc00xa,0x00x33,0x00x39,0x00x5,0x00x4,0x00x2f,0x00x35,0xc00x12,0x00x16,0x00xa,0x00x2,0x10x0,0x00xde,0x00x0,0x00x9,0x00x7,0x00x0,0x40x74,0x650x73,0x740x0,0x50x0,0x50x1,0x00x0,0x00x0,0x00xa,0x00xa,0x00x8,0x00x1d,0x00x17,0x00x18,0x00x19,0x00xb,0x00x2,0x10x0,0x00x23,0x00xa3,0x10x2,0x30x4,0x50x6,0x70x8,0x90xa,0xb0xc,0xd0xe,0xf0x10,0xca0xd3,0x190x2a,0xd40x2,0x1b0x9a,0x7b0x92,0xf70xd3,0x930x31,0x800xa4,0x300x61,0x20x1,0x10x2,0x20x3,0x420x4,0x20xc0,0x130x4,0x00x4,0x300xa3,0xfe0x4f,0x4c0xe,0xa70xb4,0x80xbc,0x900x6,0x110x50,0xb00xc1,0x540xd3,0x8b0xb2,0xfe0x45,0xfd0xd4,0x2d0xe2,0x750x25,0xad0x75,0x9e0x3c,0xab0x64,0x5a0x99,0x250x8b,0xc20xf3,0xb20xc9,0xdf0x31,0x3b0xad,0x9b0x49,0xd70xa1,0x40x2,0x20x4,0xd20xa2,0x40x2,0x20x1c,0x200xa4,0x20x4,0x00xa6,0x60x4,0x40x74,0x650x73,0x740xb1,0x30x1,0x10xff,0xb20x3,0x20x1,0x1d0xd9,0xe00xcb,0x990x9c,0xcc0x6d,0xa0x59,0x410xbe,0xdb0xeb,0x7e0x22,0xe90xed,0x140xd8,0x700x91,0x100x18,0xc70xdb,0xd60x85,0x870xf3,0x690x5b,0x450xff,0x10x0,0x10x0,0x80x17,0x00x0,0x00x16,0x00x0,0x160x3,0x20x0,0x250x10,0x00x0,0x210x20,0xf40xfe,0xd30xbb,0xe50xd3,0x310x38,0xf70x4e,0xed0xc4,0x7d0xe1,0x490x35,0x1f0xed,0xc10x8,0x280xd1,0x910xff,0x930xbb,0x900x4f,0x20xc0,0x5a0x4e,0x140x3,0x20x0,0x10x1,0x160x3,0x20x0,0x160x3,0x00xe2,0x30x3,0x10x2,0x30x4,0x50x6,0x70x8,0x90xa,0xb0xc,0xd0xe,0xf0x10,0x110x12,0x13",
        "libsndfile": "0x460x4f,0x520x4d,0xed0xff,0xff0x2d,0x410x49,0x460x46,0x430x4f,0x4d0x4d,0x00x0,0x00x12,0x00x2,0x00x0,0x00x0,0x00x10,0x520x4d,0xed0xff,0xff0x2d,0x410x49,0x460x46,0x430x4f,0x4d0x4d,0x00x0,0x00x12,0x00x2,0x00x0,0x00x0,0x00x10,0x400x4d,0x00x0,0x10xe8,0x3c0x6a,0x460x43,0x430x4f,0x4d0x4d,0x00x0,0x00x26,0x00x2,0x00x0,0x00xa9,0x00x10,0x400xf1,0xab0x44,0x00xcb,0x00x0,0x00x49,0x4d0x0,0x00x1,0xe80x41,0xae0x0,0x100x40,0xf10xab,0x440x0,0x00x0,0x00x0,0x10x69,0x6d0x61",
        "kamailio parse_msg": "0x670x6f,0x640x76,0x710x6e,0x780x66,0x6a0x69,0x700xe9,0x10x1,0x10x87,0x6e0x20,0x700xef,0xe10x83,0x00x1c,0x10x6c,0x200x53,0x590x5a,0x00x0,0x00x0,0x00xa,0x2d0x2d,0x6a0x6d,0x740x74,0x3a0x2c,0x2e0x2e,0x2c0xf6,0x920x16,0x10x33,0x350x4f,0x10xdd,0x5c0x59,0x10x1,0x820x2c,0x200x46,0xf0x20,0x00x4b,0x510x33,0x200x20,0x200x20,0x200x20,0x170x87,0x200x4b,0x4f0x35,0xa0x6d,0x3a0x2c,0x2c0x2c,0x2c0xee,0x7a0x79,0x630x6b,0x760x9,0xd0x4b,0x520x70,0x4f0x34,0x690xbf,0xbf0xcd,0x7f0x70,0x90x9,0x90x9,0x10x1,0x520x4a,0x320xa,0x560x3a,0x90x9,0x90x9,0x90x9,0x530x49,0x500x2f,0x320x2e,0x300x20,0x90x2f,0x770x77,0x770x77,0x770x77,0x770x77,0x770x77,0x770x77,0x770x77,0x770x77,0x770x77,0x770x77,0x770x77,0x770x77,0x770x77,0x770x77,0x770x77,0x770x77,0x200x20,0x200x20,0x200x52,0x470x34,0x670x6f,0x640x76,0x710x6e,0x740x3a,0x2d0x2e,0x2d0x2e,0x2e0x2d,0x6a0x73,0x670x74,0x3a0xa,0xa0xf,0x10x58,0x510x8,0x00x4e,0x770x56,0x00xa2,0x5e0x56,0x00x0,0xfd0x9,0x90x9,0x90x9,0x90x89,0x6f0x17,0x540x4c,0x390xa,0x740x3a,0x2d0x6b,0x720x79,0x200x62,0x790x6f,0x770x71,0x6b0x0,0x870x61,0x2c0x2c,0xfd0x9,0x90x9,0x90x68,0x69",
        "php exif": "0x490x49,0x2a0x0,0x80x0,0x00x0,0x50x0,0x9f0x9c,0x10x0,0x160x0,0x00x0,0x4a0x0,0x00x0,0x9e0x9c,0x10x0,0x180x0,0x00x0,0x600x0,0x00x0,0x9f0x9c,0x10x0,0x140x0,0x00x0,0x780x0,0x00x0,0x9c0x9c,0x10x0,0x3c0x0,0x00x0,0x8c0x0,0x00x0,0x9f0x9c,0x10x0,0x120x0,0x00x0,0xc80x0,0x00x0,0x00x0,0x00x0,0x100x8d,0x800xf2,0x2f0x7f,0x00x0,0x20x0,0x100x0,0x40x0,0x00x0,0x00x0,0x00x0,0x20x0,0x00x0,0x00x40,0xc90x46,0xd90x1,0x750x0,0x620x0,0x6a0x0,0x650x0,0x00x0,0x7f0xff,0xd80xff,0xe00x0,0x100x4a,0x460x49,0x460x0,0x10x2,0x10x0,0x480x0,0x480x0,0x00xff,0xe10x0,0xe20x45,0x780x69,0x660x0,0x00x49,0x490x2a,0x00x8,0x00x0,0x00x5,0x00x9f,0x9c0x1,0x00x16,0x00x0,0x00x4a,0x00x0,0x00x9e,0x9c0x1,0x00x18,0x00x0,0x00x60,0x00x0,0x00x9d,0x9c0x1,0x00x14,0x00x0,0x00x78,0x00x0,0x00x9c,0x9c0x1,0x00x3c,0x00x0,0x00x8c,0x00x0,0x00x9b,0x00x0,0x00xc8,0x00x0,0x00x0,0x00x0,0xff0xd8,0x300x0,0x100x30",
        "php parser": """<?php

class Foo {
    function testParentClass(Foo $foo) {}
    function testBothClass(Foo $foo) {}
    function testChildClass($foo) {}
    function testNoneClass($foo) {}
}

class Bar extends Foo {
    function	xclyIcjdgvAlass($foo) {}
    function testBothClass(Foo $foo) {}
    function testChildClass(Foo $foo) {}
    function testNoneClass($foo) {}
}

?>""",
        "php unserialize": """O:8:""00011110""?""",
        "json": "[0e]",
        "libxml": """<?xml version=""1.0""?>
<!DOCTYPE EXAMPLE SYSTEM ""a%aq"" [
<!ENTITY xml ""AxzfyhoagF>
]>
<EXAMvLE>
  &title;
  Ytzl pdwm wj eertW OVO- wgd &xml; ald lixl nj ds epdfafeb <WQU hybe Lblqql Mftqpdob"">
<!ENTITY title PUBLIC ""urn:publicid:%W"" ""mjsvc.jx"">
<!ENTITY uhfge SYSTEM ""kjx.yv_"" NDATA GIF>
]>
<EXAMvLE>
  &title;
  Ytzl pdwm wj eertW OVO- wgd &xml; ald lixl nj ds epdfafeb <WQU hyb=""tmemc""/>
</EXAMPLE>""",
        "sqlite3": """ 	  SELECT @ilM data LIMIT 400;
  CREATE INrsfror(6,0)-	randomblob(5/0)""",
        "lua": """if b then break elsjjjj= a + 1 end
     if b then break elsjjjj= a + 1 end
        end     end""",
    }

    seed_queue = [example[args.fuzzing_target]]

    generation_kwargs = {
        "do_sample": True,
        "min_length": -1,
        "top_p": 0.92,  # 0.9
        "top_k": 50,
        "temperature": args.temperature,
        "pad_token_id": tokenizer.bos_token_id,
    }

    while True:
        global seeds_from_fuzzer
        is_from_fuzzer = False
        current_seed = random.choice(seed_queue)

        if seeds_from_fuzzer:
            current_seed = seeds_from_fuzzer.pop()
            if len(seed_queue) > 30:
                seed_queue = []
            seed_queue.append(current_seed)
            is_from_fuzzer = True

        formatted_chunks = []
        if not args.if_text:
            for i in range(0, len(current_seed), 4):
                if i + 3 < len(current_seed):
                    formatted_chunks.append(
                        f"0x{current_seed[i:i+2]}0x{current_seed[i+2:i+4]}"
                    )
                else:
                    # If no pair, add the single element
                    formatted_chunks.append(f"0x{current_seed[i:]}")
        else:
            formatted_chunks.append(current_seed)
        prompt = (
            "### Input: ```Based on below hex "
            + args.fuzzing_target
            + " seed, mutate a new "
            + args.fuzzing_target
            + " seed. Make sure the example is complete and valid. "
            + ",".join(formatted_chunks)
            + "```"
        )

        query_tensors = tokenizer(prompt, return_tensors="pt")["input_ids"].to("cuda")
        response_tensors = model.generate(
            input_ids=query_tensors,
            max_new_tokens=500,
            **generation_kwargs,
        )

        response = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)
        # Compute sentiment score
        global uid, seed_id_map, id_rwd_map, message_queue
        for r in response:
            seed = hex_string_to_hex(r, args.fuzzing_target, args.if_text)
            seed_id_map[seed] = uid + os.getpid()
            # id_rwd_map[uid + os.getpid()] = float(0.0)
            message_queue.append(seed)
            if is_from_fuzzer:
                print("sff:::", seed[:15])
            else:
                print("seed:::", seed[:15])
        uid += 8
        torch.cuda.empty_cache()


if __name__ == "__main__":
    t = threading.Thread(
        target=mq_thread,
        args=(),
    )
    t.start()
    main()