Virtualize asm 코드 분할 및 라벨링 자동화

카테고리 없음

Virtualize asm 코드 분할 및 라벨링 자동화

snejs 2025. 9. 24. 20:24

코드 생성에는 GPT를 활용하였다.

문제 상황

Virtualize 난독화가 적용된 asm 코드를 입력으로 하여 코드를 분할하고 라벨링을 자동화 할 수 있는 코드 생성

가정

입력 : Virtualize 난독화가 적용된 asm 코드
출력 : 분할된 asm 코드 집합 + labeling된 데이터
요구사항 : "call"을 기준으로 코드 분할

위 상황에 대해 구체적인 문제에 대입 가능하도록 하는 코드 초안을 생성

사용된 원본 .c 코드는 약 4KB 코드를 사용하였으며, Virtualize 난독화를 적용한 후에는 약 695KB의 코드가 되었다.

Python Code 생성

직접 pass를 사용하고 생성하는 것을 바로 하기에는 힘들 것으로 보여 python code를 1차적으로 생성하고,

요구사항에 맞춰서 수정하였다.

python code 에 사용된 입력은 clang -S 으로 컴파일된 .s 파일을 사용하였다.

정규식을 이용하여 단순 "call" 찾아 코드를 분할 한 뒤, .jsonl에 id와 labeling 등의 정보를 저장할 수 있게 했다.

import re
import json
import os

def split_asm_file(input_file, output_file):
    # 폴더 & 기본 이름 준비
    base_name = os.path.splitext(os.path.basename(input_file))[0]
    folder = os.path.join(os.path.dirname(output_file) or ".", base_name)
    os.makedirs(folder, exist_ok=True)

    with open(input_file, "r", encoding="utf-8") as f:
        lines = f.readlines()

    chunks_meta = []
    current_chunk = []
    chunk_id = 0

    for line in lines:
        line_strip = line.strip()
        if not line_strip:
            continue  # 빈 줄 무시

        current_chunk.append(line_strip)

        # 함수 시작 라벨 추출(.globl, <func>: 같은 부분)
        if re.match(r"^[a-zA-Z0-9_]+:$", line_strip):
            labels = ["FUNC_ENTRY"]
        else:
            labels = []

        # call 명령어 감지
        if re.search(r"\bcall\b", line_strip):
            labels.append("CALL_SITE")

            # 청크 id 및 파일명 생성
            chunk_name = f"{base_name}_chunk_{chunk_id}"
            chunk_filename = chunk_name + ".s"
            chunk_path = os.path.join(folder, chunk_filename)

            # 청크 내용 파일로 저장 (각 라인 뒤에 newline 추가)
            with open(chunk_path, "w", encoding="utf-8") as cf:
                cf.write("\n".join(current_chunk) + "\n")

            # 메타 정보만 JSONL에 쓸 리스트에 추가 (asm 필드 없음)
            chunks_meta.append({
                "id": chunk_name,
                "labels": labels,
                "file": os.path.relpath(chunk_path)  # 상대경로로 저장 
            })

            chunk_id += 1
            current_chunk = []  # 새 청크 시작

    # 마지막 청크가 남아있다면 파일로 저장하고 메타에 추가
    if current_chunk:
        chunk_name = f"{base_name}_chunk_{chunk_id}"
        chunk_filename = chunk_name + ".s"
        chunk_path = os.path.join(folder, chunk_filename)

        with open(chunk_path, "w", encoding="utf-8") as cf:
            cf.write("\n".join(current_chunk) + "\n")

        chunks_meta.append({
            "id": chunk_name,
            "labels": [],
            "file": os.path.relpath(chunk_path)
        })

    # JSONL 저장 (asm 필드 없음)
    with open(output_file, "w", encoding="utf-8") as f:
        for meta in chunks_meta:
            f.write(json.dumps(meta, ensure_ascii=False) + "\n")

    print(f"[+] Saved {len(chunks_meta)} chunks to {output_file}")
    print(f"[+] Chunk files are in folder: {folder}")

if __name__ == "__main__":
    input_file = "test_out.s"  # 분석할 asm 파일
    output_file = "test_out.jsonl"  # 결과 저장할 JSONL 파일

    split_asm_file(input_file, output_file)

결과 예시

생성된 파일은 약 1KB ~ 300KB 의 파일들이 41개 생성되었다.

C++ Code 생성

위에서 생성한 python 코드를 바탕으로 LLVM Pass 를 활용한 C++ 코드를 생성하도록 했다.

python에서 불필요한 정보들을 제거하고, 하나의 코드에서 분할된 코드 집합과 jsonl 파일을 한 파일에 생성되도록 수정하였다.

// CallSplitterPass.cpp
#include "llvm/ADT/StringRef.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PassManager.h"
#include "llvm/Passes/PassPlugin.h"
#include "llvm/Passes/PassBuilder.h" 
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/Path.h"
#include "llvm/Support/raw_ostream.h"

#include <filesystem>   
#include <vector>
#include <string>
#include <fstream>

using namespace llvm;
namespace fs = std::filesystem;

static cl::opt<std::string> OutputDir("split-asm-dir",
    cl::desc("Directory to place chunk files"),
    cl::init("split_asm_out"));

static cl::opt<std::string> JsonOut("split-asm-json",
    cl::desc("JSONL output file"),
    cl::init("chunks.jsonl"));

namespace {

struct CallSplitterPass : public PassInfoMixin<CallSplitterPass> {
    PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
        // 1) 모듈 이름 가져오기
        std::string moduleName = M.getSourceFileName();
        if (moduleName.empty()) {
            StringRef mid = M.getModuleIdentifier();
            moduleName = sys::path::stem(mid).str(); // stem(StringRef) -> StringRef, .str()로 std::string으로 변환
        }
        if (moduleName.empty())
            moduleName = "module";

        // 2) 출력 디렉토리
        fs::path baseDir = OutputDir.getValue();
        std::error_code ec_fs;
        
        fs::create_directories(baseDir, ec_fs);
        if (ec_fs) {
            errs() << "Warning: cannot create dir '" << baseDir.string() << "': " << ec_fs.message() << "\n";
        }

        fs::path moduleDir = baseDir / moduleName;

        fs::create_directories(moduleDir, ec_fs);
        if (ec_fs) {
            errs() << "Warning: cannot create module dir '" << moduleDir.string() << "': " << ec_fs.message() << "\n";
        }

        fs::path jsonPath = moduleDir / JsonOut.getValue();
        // 3) JSONL 파일 (덮어쓰기)
        std::ofstream jsonl(jsonPath, std::ios::out | std::ios::trunc);
        if (!jsonl.is_open()) {
            errs() << "Error: cannot open JSONL file: " << JsonOut << "\n";
            return PreservedAnalyses::all();
        }

        unsigned chunk_id = 0;

        for (Function &F : M) {
            if (F.isDeclaration()) continue;

            std::vector<std::string> currentChunkLines;
            std::vector<std::string> currentLabels;
            bool seenAnyInstInFunction = false;

            for (BasicBlock &BB : F) {
                for (Instruction &I : BB) {
                  // 라벨 추가
                    if (!seenAnyInstInFunction) {
                        currentLabels.push_back("FUNC_ENTRY");
                        seenAnyInstInFunction = true;
                    }

                    // instruction -> 문자열
                    std::string instStr;
                    raw_string_ostream rso(instStr);
                    I.print(rso);
                    rso.flush();
                    currentChunkLines.push_back(instStr);

                    // Call 계열이면 청크 저장
                    if (isa<CallBase>(&I)) {
                        currentLabels.push_back("CALL_SITE");

                        std::string chunkName = "chunk_" + std::to_string(chunk_id);
                        std::string chunkFilename = chunkName + ".s";
                        fs::path chunkPath = moduleDir / chunkFilename;

                        // 파일 쓰기 (llvm raw_fd_ostream 사용)
                        std::error_code ec;
                        raw_fd_ostream ofs(chunkPath.string(), ec, sys::fs::OF_Text);
                        if (ec) {
                            errs() << "Warning: cannot write chunk file " << chunkPath.string() << ": " << ec.message() << "\n";
                        } else {
                            for (auto &L : currentChunkLines) ofs << L << "\n";
                        }

                        // JSONL 기록 
                        jsonl << "{\"id\":\"" << chunkName << "\",";
                        jsonl << "\"labels\":[";
                        for (size_t li = 0; li < currentLabels.size(); ++li) {
                            jsonl << "\"" << currentLabels[li] << "\"";
                            if (li + 1 < currentLabels.size()) jsonl << ",";
                        }
                        jsonl << "]";
                        jsonl << "}\n";

                        ++chunk_id;
                        currentChunkLines.clear();
                        currentLabels.clear();
                        seenAnyInstInFunction = false; // 다음 청크는 새로운 시작
                    }
                }
            }

            // 함수 끝에 남은 청크 저장
            if (!currentChunkLines.empty()) {
                std::string chunkName = "chunk_" + std::to_string(chunk_id);
                std::string chunkFilename = chunkName + ".s";
                fs::path chunkPath = moduleDir / chunkFilename;

                std::error_code ec;
                raw_fd_ostream ofs(chunkPath.string(), ec, sys::fs::OF_Text);
                if (ec) {
                    errs() << "Warning: cannot write chunk file " << chunkPath.string() << ": " << ec.message() << "\n";
                } else {
                    for (auto &L : currentChunkLines) ofs << L << "\n";
                }

                jsonl << "{\"id\":\"" << chunkName << "\",";
                jsonl << "\"labels\":[";
                for (size_t li = 0; li < currentLabels.size(); ++li) {
                    jsonl << "\"" << currentLabels[li] << "\"";
                    if (li + 1 < currentLabels.size()) jsonl << ",";
                }
                jsonl << "]";
                jsonl << "}\n";

                ++chunk_id;
            }
        }

        jsonl.close();
        errs() << "[CallSplitterPass] wrote " << chunk_id << " chunks to directory '" << moduleDir.string()
               << "' and JSONL '" << JsonOut << "'\n";

        return PreservedAnalyses::all();
    }
};

} // namespace

// plugin registration
extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo llvmGetPassPluginInfo() {
    return {
        LLVM_PLUGIN_API_VERSION, "CallSplitterPass", LLVM_VERSION_STRING,
        [](PassBuilder &PB) {
            PB.registerPipelineParsingCallback(
                [](StringRef Name, ModulePassManager &MPM,
                   ArrayRef<PassBuilder::PipelineElement>) -> bool {
                    if (Name == "call-splitter") {
                        MPM.addPass(CallSplitterPass());
                        return true;
                    }
                    return false;
                });
        }
    };
}

사용한 컴파일 코드

clang++ -std=c++17 -fPIC -shared CallSplitterPass.cpp -o CallSplitterPass.so \
  `llvm-config --cxxflags`

clang -O1 -emit-llvm -c -Xclang -disable-O0-optnone -o test_out.bc test_out.c

opt -load-pass-plugin=./CallSplitterPass.so -passes="call-splitter" \
    -split-asm-dir=split_asm_out -split-asm-json=chunks.jsonl test_out.bc -o /dev/null