File size: 5,355 Bytes
c01b9cd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
import fs from "fs/promises";
import path from "path";

export interface ProjectContext {
  remappings: string;       // Content of remappings.txt or foundry.toml [profile.default.remappings]
  existingTestImports: string; // First few import lines from an existing test file
  foundryTomlProfile: string; // Relevant foundry.toml settings (src, libs)
  existingTestFilePath: string | null; // Relative path to an existing test file for reference
}

/**
 * Extracts project-level context needed for correct import paths in PoC tests.
 * Reads remappings.txt, foundry.toml, and the first existing test file in the project.
 */
export async function extractProjectContext(sandboxDir: string): Promise<ProjectContext> {
  let remappings = "";
  let foundryTomlProfile = "";
  let existingTestImports = "";
  let existingTestFilePath: string | null = null;

  // 1. Read remappings.txt
  try {
    const remappingsPath = path.join(sandboxDir, "remappings.txt");
    remappings = await fs.readFile(remappingsPath, "utf-8");
  } catch {
    // fallback: try to extract from foundry.toml
  }

  // 2. Read foundry.toml for additional context
  try {
    const foundryTomlPath = path.join(sandboxDir, "foundry.toml");
    const tomlContent = await fs.readFile(foundryTomlPath, "utf-8");
    // Extract relevant lines (src, libs, remappings)
    const relevantLines = tomlContent
      .split("\n")
      .filter(l =>
        l.includes("src") ||
        l.includes("libs") ||
        l.includes("remapping") ||
        l.includes("[profile")
      )
      .slice(0, 20)
      .join("\n");
    foundryTomlProfile = relevantLines;

    // If no remappings.txt, try to extract from foundry.toml remappings array
    if (!remappings) {
      const remappingMatch = tomlContent.match(/remappings\s*=\s*\[([\s\S]*?)\]/);
      if (remappingMatch) {
        remappings = remappingMatch[1]
          .split(",")
          .map(s => s.trim().replace(/^["']|["']$/g, ""))
          .filter(Boolean)
          .join("\n");
      }
    }
  } catch {
    // ignore
  }

  // 3. Find an existing test file to use as import reference
  try {
    const testDir = path.join(sandboxDir, "test");
    const testFile = await findFirstTestFile(testDir);
    if (testFile) {
      existingTestFilePath = path.relative(sandboxDir, testFile);
      const testContent = await fs.readFile(testFile, "utf-8");
      // Extract the first 20 lines which typically contain imports
      existingTestImports = testContent
        .split("\n")
        .slice(0, 25)
        .filter(l => l.startsWith("import") || l.startsWith("pragma") || l.startsWith("//") || l.startsWith("contract") || l.startsWith("abstract"))
        .join("\n");
    }
  } catch {
    // ignore
  }

  return { remappings, existingTestImports, foundryTomlProfile, existingTestFilePath };
}

/**
 * Recursively finds the best reference test file in the test directory.
 * Strategy:
 *  1. Collect ALL .t.sol files (excluding Exploit.t.sol) across all subdirs (BFS)
 *  2. Pick the one with the most import lines (most context-rich)
 *  3. Fallback to .sol files that have "import" statements (e.g. BaseTest.sol, Fixture.sol)
 *  4. Skip pure mock contracts (files in "mock" directories or named *Mock.sol)
 */
async function findFirstTestFile(dir: string): Promise<string | null> {
  const tSolFiles: string[] = [];
  const solFiles: string[] = [];

  // BFS collect all files
  const queue = [dir];
  let depth = 0;
  while (queue.length > 0 && depth < 4) {
    const currentDepth: string[] = [...queue];
    queue.length = 0;
    depth++;
    for (const currentDir of currentDepth) {
      let entries: import('fs').Dirent[];
      try {
        entries = await fs.readdir(currentDir, { withFileTypes: true }) as import('fs').Dirent[];
      } catch {
        continue;
      }
      for (const entry of entries) {
        const fullPath = path.join(currentDir, entry.name);
        const name = entry.name as string;
        if (entry.isFile()) {
          if (name === "Exploit.t.sol") continue; // skip our own file
          if (name.endsWith(".t.sol")) {
            tSolFiles.push(fullPath);
          } else if (name.endsWith(".sol")) {
            // Skip mock contracts
            const isMock = name.toLowerCase().includes("mock") || currentDir.toLowerCase().includes("mock");
            if (!isMock) {
              solFiles.push(fullPath);
            }
          }
        } else if (entry.isDirectory()) {
          queue.push(fullPath);
        }
      }
    }
  }

  // Pick the .t.sol file with the most import lines (richest context)
  if (tSolFiles.length > 0) {
    let bestFile = tSolFiles[0];
    let bestImportCount = 0;
    for (const f of tSolFiles.slice(0, 10)) { // check up to 10
      try {
        const content = await fs.readFile(f, "utf-8");
        const importCount = (content.match(/^import/gm) ?? []).length;
        if (importCount > bestImportCount) {
          bestImportCount = importCount;
          bestFile = f;
        }
      } catch { /* skip */ }
    }
    return bestFile;
  }

  // Fallback: .sol files that have import statements (like BaseTest.sol, Fixture.sol)
  for (const f of solFiles.slice(0, 10)) {
    try {
      const content = await fs.readFile(f, "utf-8");
      if (content.includes("import ")) return f;
    } catch { /* skip */ }
  }

  return null;
}