Split config and ability to change api key

packages/ethernaut-ai-ui/src/index.js

@@ -4,6 +4,6 @@ require('ethernaut-ui/src/index')
 require('ethernaut-ai/src/index')
 
 extendEnvironment((hre) => {
-  const config = hre.scopes.ai.tasks.config
+  const config = hre.scopes.ai.tasks.model
   config.paramDefinitions.model.prompt = require('./prompts/model')
 })

packages/ethernaut-ai/src/tasks/key.js

@@ -0,0 +1,31 @@
+const types = require('ethernaut-common/src/validation/types')
+const output = require('ethernaut-common/src/ui/output')
+const storage = require('ethernaut-common/src/io/storage')
+const { setEnvVar } = require('ethernaut-common/src/io/env')
+
+require('../scopes/ai')
+  .task('key', 'Sets the openai api key')
+  .addParam('apiKey', 'The openai api key to use', undefined, types.string)
+  .setAction(async ({ apiKey }, hre) => {
+    try {
+      const config = storage.readConfig()
+
+      let summary = []
+
+      if (apiKey) {
+        const currentKey = process.env.OPENAI_API_KEY
+        setEnvVar('OPENAI_API_KEY', apiKey)
+        summary.push(`- API Key set to ${apiKey} (was ${currentKey})`)
+      }
+
+      storage.saveConfig(config)
+
+      if (summary.length === 0) {
+        summary.push('No changes')
+      }
+
+      return output.resultBox(summary.join('\n'))
+    } catch (err) {
+      return output.errorBox(err)
+    }
+  })

packages/ethernaut-ai/src/tasks/config.js → packages/ethernaut-ai/src/tasks/model.js

@@ -3,7 +3,7 @@ const output = require('ethernaut-common/src/ui/output')
 const storage = require('ethernaut-common/src/io/storage')
 
 require('../scopes/ai')
-  .task('config', 'Configures ai scope parameters')
+  .task('model', 'Sets the openai model')
   .addParam('model', 'The openai model to use', undefined, types.string)
   .setAction(async ({ model }, hre) => {
     try {

packages/ethernaut-common/src/io/env.js

@@ -8,35 +8,57 @@ const { getEthernautFolderPath } = require('ethernaut-common/src/io/storage')
 const envPath = path.join(getEthernautFolderPath(), '.env')
 
 function refreshEnv() {
+  // Create the .env file if it doesn't exist
+  if (!fs.existsSync(envPath)) {
+    debug.log('No .env file found, creating one...', 'env')
+    fs.writeFileSync(envPath, '')
+  }
+
+  // Load the .env file
   require('dotenv').config({ path: envPath })
 }
 
 async function checkEnvVar(varName, message) {
+  // Check if the env var exists at runtime
   if (process.env[varName]) {
     debug.log(`Environment variable ${varName} found`, 'env')
     return
   }
-
   debug.log(`Environment variable ${varName} not found - collecting it...`)
 
-  if (!fs.existsSync(envPath)) {
-    debug.log('No .env file found, creating one...', 'env')
-    fs.writeFileSync(envPath, '')
-  }
+  // Collect the env var from the user
+  const varValue = await prompt({
+    type: 'input',
+    message: `Please provide a value for ${varName}${message ? `. ${message}` : ''}`,
+  })
+
+  // Save the env var to the .env file
+  setEnvVar(varName, varValue)
+}
 
-  const envConfig = dotenv.parse(fs.readFileSync(envPath))
-  if (!envConfig[varName]) {
-    const varValue = await prompt({
-      type: 'input',
-      message: `Please provide a value for ${varName}${message ? `. ${message}` : ''}`,
-    })
-    debug.log(`Saved environment variable ${varName} in .env file...`)
-    fs.appendFileSync(envPath, `${varName}=${varValue}`)
-    require('dotenv').config({ path: envPath })
+function loadEnvConfig() {
+  return dotenv.parse(fs.readFileSync(envPath))
+}
+
+function setEnvVar(varName, varValue) {
+  // Load and set the env var
+  const envConfig = loadEnvConfig()
+  envConfig[varName] = varValue
+
+  // Write the env var to the .env file
+  let envFileContent = ''
+  for (const [key, value] of Object.entries(envConfig)) {
+    envFileContent += `${key}=${value}\n`
   }
+  fs.writeFileSync(envPath, envFileContent)
+  debug.log(`Saved environment variable ${varName} in .env file...`)
+
+  process.env[varName] = varValue
+  refreshEnv()
 }
 
 module.exports = {
   refreshEnv,
   checkEnvVar,
+  setEnvVar,
 }