LLM4Decompileを使う

# 要約
- wsl上でllm4decompileを動作させた
- RTX 3060 12GBで動作確認できたのはllm4decompile-6.7b-v2がギリギリ動作
- v1.5モデルではアセンブリ→cに v2モデルではデコンパイラ→cに変換する
    - v2を試した
# LLM4Decompileとは
https://github.com/albertan017/LLM4Decompile/tree/main
生成AIベースのデコンパイラ。
学習にはCをO0~3でコンパイルしたelf x86_64を使用。
# 検証環境
- NVIDIA RTX 3060 12GB
- WSL Ubuntu
- RAM 64GB

# 環境構築
vllmを使用します。wslでないと動作しないので注意
他の方法としてはDockerを使うなど
## wsl上にcudaドライバを入れる
https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&Distribution=WSL-Ubuntu&target_version=2.0

## リポジトリのクローン
ここから先は任意のディレクトリで操作してください。
クローンしなくても動くとは思う。
```bash
git clone https://github.com/albertan017/LLM4Decompile
```

## Pythonのセットアップ
README.mdだとcondaを使っているようだが、venvでも動いた。
```
python -m venv llm4dec
source llm4dec/bin/activate
```

依存関係をインストール
```
cd LLM4Decompile
pip install -r requirements.txt
```

## Ghidraのインストール
すでにある人はスキップでも化
今回はGhidra 11.1.1を使用。
ついでにjavaも入れる。11.1までjdk17。11.2からはjdk21なので注意。
```bash
sudo apt install openjdk-17-jdk
```
これで大体の準備が終わった。

# LLM4Decompileを使う
LLM4Decompile/ghidraにあるスクリプトを実行するとサンプルのコードを試せる。
6.7Bのモデルを指定しているのでVRAMが少ない場合はllm4decompile-1.3b-v2に変更する or google colaboratoryを使う。
あとghidra_path(特にバージョン)を修正すること。
demo.pyではLLM4Decompile/samples/sample.cを各最適化オプションでコンパイルしてGhidraをHeadlessモードで起動し、decompile.pyを実行しデコンパイル結果を一時的なファイルに保存する。
```bash
cd ghidra
python demo.py
```
```c
pseudo function:
# This is the assembly code:
undefined8 func0(float param_1,long param_2,int param_3)

{
  int local_10;
  int local_c;
  
  local_10 = 0;
  do {
    local_c = local_10;
    if (param_3 <= local_10) {
      return 0;
    }
    while (local_c = local_c + 1, local_c < param_3) {
      if ((float)(DAT_001020d0 &
                 (uint)(*(float *)(param_2 + (long)local_10 * 4) -
                       *(float *)(param_2 + (long)local_c * 4))) < param_1) {
        return 1;
      }
    }
    local_10 = local_10 + 1;
  } while( true );
}
# What is the source code?

refined function:
int func0(float *x, int n, float eps)
{
    int i, j;
    for (i = 0; i < n; i++)
        for (j = i + 1; j < n; j++)
            if (fabs(x[i] - x[j]) < eps)
                return 1;
    return 0;
}
```
見やすいですね

## 特定のファイルをデコンパイルする
ここらへんは正直おまけで困ったらLLM4Decompile/ghidra/demo.pyを見ながらコードを改造すれば大体やりたいことはできると思う。

executable_pathとtarget_funcを書き換えればおｋ。
すでにGhidraに読み込ませてある実行ファイルを使いたい場合はHeadlessモードのコマンドを変えればおｋ。
```python
import tempfile
import os
import subprocess
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

timeout_duration = 10

ghidra_path = "./ghidra_11.1.1_PUBLIC/support/analyzeHeadless"#path to the headless analyzer, change the path accordingly
postscript = "./decompile.py"#path to the decompiler helper function, change the path accordingly
project_path = "."#path to temp folder for analysis, change the path accordingly
project_name = "tmp_ghidra_proj"
fileName = "dec"
executable_path = "../a.out"
target_func = "main"
# model_path = 'LLM4Binary/llm4decompile-9b-v2' # V2 Model
# model_path = 'LLM4Binary/llm4decompile-6.7b-v2' # V2 Model
model_path = 'LLM4Binary/llm4decompile-1.3b-v2' # V2 Model

with tempfile.TemporaryDirectory() as temp_dir:
    pid = os.getpid()
    asm_all = {}
    output_path = os.path.join(temp_dir, f"{pid}_.c")
    command = [
        ghidra_path,
        temp_dir,
        project_name,
        "-import", executable_path,
        "-postScript", postscript, output_path,
        "-deleteProject",  # WARNING: This will delete the project after analysis
    ]
    result = subprocess.run(command, text=True, capture_output=True, check=True)
    with open(output_path,'r') as f:
        c_decompile = f.read()
    c_func = []
    flag = 0
    for line in c_decompile.split('\n'):
        if f"Function: {target_func}" in line:#**Replace** func0 with the function name you want to decompile.
            flag = 1
            c_func.append(line)
            continue
        if flag:
            if '// Function:' in line:
                if len(c_func) > 1:
                    break
            c_func.append(line)
    if flag == 0:
        raise ValueError('bad case no function found')
    for idx_tmp in range(1,len(c_func)):##########remove the comments
        if target_func in c_func[idx_tmp]:
            break
    c_func = c_func[idx_tmp:]
    input_asm = '\n'.join(c_func).strip()

    before = f"# This is the assembly code:\n"#prompt
    after = "\n# What is the source code?\n"#prompt
    input_asm_prompt = before+input_asm.strip()+after
    with open(fileName + '.pseudo','w',encoding='utf-8') as f:
        f.write(input_asm_prompt)



tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16).cuda()

with open(fileName + '.pseudo','r') as f:#optimization level O0
    asm_func = f.read()
inputs = tokenizer(asm_func, return_tensors="pt").to(model.device)
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=2048)### max length to 4096, max new tokens should be below the range
c_func_decompile = tokenizer.decode(outputs[0][len(inputs[0]):-1])

# with open(fileName + '.pseudo','r') as f:#original file
#     func = f.read()

# print(f'pseudo function:\n{func}')# Note we only decompile one function, where the original file may contain multiple functions
print(f'refined function:\n{c_func_decompile}')
```

# 検証
元のコード
これをgcc -O2でコンパイルした。
```c
#include <stdio.h>
#include <stdlib.h>
#include <time.h>

int gen_rand() {
    return rand() % 100;
}

int main() {
    int *a = (int *)malloc(3000 * 3000 * sizeof(int));
    int *b = (int *)malloc(3000 * 3000 * sizeof(int));
    int *c = (int *)calloc(3000 * 3000, sizeof(int));
    
    srand(0);
    for (int i = 0; i < 3000 * 3000; i++) {
        a[i] = gen_rand();
        b[i] = gen_rand();
    }

    clock_t start_clock, end_clock;
    start_clock = clock();
    for (int i = 0; i < 3000; i++) {
        for (int j = 0; j < 3000; j++) {
            for (int k = 0; k < 3000; k++) {
                c[j * 3000 + k] += a[j * 3000 + i] * b[i * 3000 + k];
            }
        }
    }
    end_clock = clock();
    printf("Clock: %f\n", (double)(end_clock - start_clock) / CLOCKS_PER_SEC);

    free(a);
    free(b);
    free(c);

    return 0;
}
```
## Before
```c
undefined8 main(void)

{
  ulong *puVar1;
  int *piVar2;
  int *piVar3;
  uint uVar4;
  undefined auVar5 [16];
  ulong uVar6;
  ulong uVar7;
  int iVar8;
  int iVar9;
  ulong uVar10;
  int iVar11;
  uint *__ptr;
  void *__ptr_00;
  void *__ptr_01;
  clock_t cVar12;
  long lVar13;
  clock_t cVar14;
  void *pvVar15;
  uint *puVar16;
  uint *puVar17;
  long lVar18;
  
  lVar18 = 0;
  __ptr = (uint *)malloc(36000000);
  __ptr_00 = malloc(36000000);
  __ptr_01 = calloc(9000000,4);
  srand(0);
  do {
    iVar11 = rand();
    *(int *)((long)__ptr + lVar18) = iVar11 % 100;
    iVar11 = rand();
    *(int *)((long)__ptr_00 + lVar18) = iVar11 % 100;
    lVar18 = lVar18 + 4;
  } while (lVar18 != 36000000);
  cVar12 = clock();
  lVar18 = 0;
  puVar17 = __ptr;
  do {
    pvVar15 = __ptr_01;
    puVar16 = puVar17;
    do {
      uVar4 = *puVar16;
      lVar13 = 0;
      auVar5._4_4_ = uVar4;
      auVar5._0_4_ = uVar4;
      auVar5._8_4_ = uVar4;
      auVar5._12_4_ = uVar4;
      do {
        puVar1 = (ulong *)((long)__ptr_00 + lVar13 + lVar18 * 12000);
        uVar6 = *puVar1;
        uVar7 = puVar1[1];
        puVar1 = (ulong *)((long)__ptr_00 + lVar13 + lVar18 * 12000);
        uVar10 = puVar1[1];
        piVar2 = (int *)((long)pvVar15 + lVar13);
        iVar11 = piVar2[1];
        iVar8 = piVar2[2];
        iVar9 = piVar2[3];
        piVar3 = (int *)((long)pvVar15 + lVar13);
        *piVar3 = *piVar2 + (int)((*puVar1 & 0xffffffff) * (ulong)uVar4);
        piVar3[1] = iVar11 + (int)((uVar6 >> 0x20) * (ulong)uVar4);
        piVar3[2] = iVar8 + (int)((uVar10 & 0xffffffff) * (ulong)uVar4);
        piVar3[3] = iVar9 + (int)((uVar7 >> 0x20) * (auVar5._8_8_ >> 0x20));
        lVar13 = lVar13 + 0x10;
      } while (lVar13 != 12000);
      puVar16 = puVar16 + 3000;
      pvVar15 = (void *)((long)pvVar15 + 12000);
    } while (puVar16 != puVar17 + 9000000);
    lVar18 = lVar18 + 1;
    puVar17 = puVar17 + 1;
  } while (lVar18 != 3000);
  cVar14 = clock();
  __printf_chk((double)(cVar14 - cVar12) / _DAT_00102010,1,"Clock: %f\n");
  free(__ptr);
  free(__ptr_00);
  free(__ptr_01);
  return 0;
}
```
## After
```c
int main(int argc, char *argv[])
{
    int i, j, k;
    int *a = malloc(3000 * 3000 * sizeof(int));
    int *b = malloc(3000 * 3000 * sizeof(int));
    int *c = calloc(3000 * 3000, sizeof(int));

    srand(0);

    for (i = 0; i < 3000 * 3000; i++) {
        a[i] = rand() % 100;
        b[i] = rand() % 100;
    }

    clock_t start = clock();

    for (i = 0; i < 3000; i++) {
        for (j = 0; j < 3000; j++) {
            for (k = 0; k < 3000; k++) {
                c[i * 3000 + j] += a[i * 3000 + k] * b[k * 3000 + j];
            }
        }
    }

    __printf_chk((double)(clock() - start) / CLOCKS_PER_SEC, 1, "Clock: %f\n");

    free(a);
    free(b);
    free(c);

    return 0;
}
```
大体あってそうに見えるけどループ内の配列のアドレスが元コードとは若干違った。
でもすごい！