"use strict";
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    var desc = Object.getOwnPropertyDescriptor(m, k);
    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
      desc = { enumerable: true, get: function() { return m[k]; } };
    }
    Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
    Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
    o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
    var ownKeys = function(o) {
        ownKeys = Object.getOwnPropertyNames || function (o) {
            var ar = [];
            for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
            return ar;
        };
        return ownKeys(o);
    };
    return function (mod) {
        if (mod && mod.__esModule) return mod;
        var result = {};
        if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
        __setModuleDefault(result, mod);
        return result;
    };
})();
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.doTrusysSecurityTest = doTrusysSecurityTest;
exports.redteamRunCommand = redteamRunCommand;
const chalk_1 = __importDefault(require("chalk"));
const dedent_1 = __importDefault(require("dedent"));
const fs = __importStar(require("fs"));
const js_yaml_1 = __importDefault(require("js-yaml"));
const lodash_1 = __importDefault(require("lodash"));
const os = __importStar(require("os"));
const path = __importStar(require("path"));
const zod_1 = require("zod");
const cliState_1 = __importDefault(require("../../cliState"));
const eval_1 = require("../../commands/eval");
const constants_1 = require("../../constants");
const envars_1 = require("../../envars");
const evaluator_1 = require("../../evaluator");
const evaluatorHelpers_1 = require("../../evaluatorHelpers");
const accounts_1 = require("../../globalConfig/accounts");
const cloud_1 = require("../../globalConfig/cloud");
const logger_1 = __importDefault(require("../../logger"));
const providers_1 = require("../../providers");
const telemetry_1 = __importDefault(require("../../telemetry"));
const util_1 = require("../../util");
const apiCall_1 = require("../../util/apiCall");
const cloud_2 = require("../../util/cloud");
const shared_1 = require("../shared");
const sharedFrontend_1 = require("../sharedFrontend");
const generate_1 = require("./generate");
const poison_1 = require("./poison");
const UUID_REGEX = /^[A-Fa-f0-9]{8}-[A-Fa-f0-9]{4}-[A-Fa-f0-9]{4}-[A-Fa-f0-9]{4}-[A-Fa-f0-9]{12}$/;
const runSecurityTest = async (reqBody, // You can replace 'any' with a more specific type if available
apiHost, apiKey) => {
    try {
        const apiUrl = `${apiHost}/api/redteam/run`;
        const result = await (0, apiCall_1.apiCall)({
            method: 'POST',
            params: {
                ...reqBody,
            },
            headers: {
                ...eval_1.DEFAULT_HEADERS,
                'x-api-key': apiKey,
            },
        }, apiUrl);
        if (result.error) {
            logger_1.default.error(`❌ Error while security run: ${result.error}`);
            return null;
        }
        return result;
    }
    catch (error) {
        logger_1.default.error(`❌ Error while security run: ${error}`);
        return null;
    }
};
const getSecurityConfigDetails = async (securityConfigName, apiKey, apiHost) => {
    try {
        const apiUrl = `${apiHost}/api/redteam/config/details`;
        const result = await (0, apiCall_1.apiCall)({
            method: 'GET',
            params: {
                name: securityConfigName,
            },
            headers: {
                ...eval_1.DEFAULT_HEADERS,
                'x-api-key': apiKey,
            },
        }, apiUrl);
        if (result.error) {
            logger_1.default.error(`❌ Error fetching security config details: ${result.error}`);
            return null;
        }
        return result;
    }
    catch (error) {
        logger_1.default.error(`❌ Error fetching security config details: ${error}`);
        return null;
    }
};
const securityTestRunHandler = async (testSuite, apiHost, apiKey) => {
    try {
        const response = await runSecurityTest(testSuite, apiHost, apiKey);
        // logger.info(`SecurityTestRun Response: ${JSON.stringify(response)}`);
        if (response?.error || !response?.id) {
            logger_1.default.error(`❌ Error running security test - ${response?.error ?? ''}`);
            return null;
        }
        const job = response;
        return new Promise((resolve) => {
            const intervalId = setInterval(async () => {
                const progressResponse = await (0, eval_1.postTestEvalWithJobId)(job.id, apiKey, apiHost);
                // logger.info(`Response from progressResponse: ${JSON.parse(progressResponse)} `);
                if (!progressResponse.status) {
                    clearInterval(intervalId);
                    logger_1.default.error(`❌ Error running test suite - ${progressResponse?.error}`);
                    // setIsRunning(false);
                    resolve({ success: false, error: progressResponse?.error });
                    return;
                }
                const progressData = progressResponse;
                if (progressData.status === 'complete') {
                    clearInterval(intervalId);
                    // setIsRunning(false);
                    // Refresh if already on the test-runs-list page
                    resolve({ success: true, testRunId: job.testRunId });
                }
                else if (progressData.status === 'failed') {
                    clearInterval(intervalId);
                    // setIsRunning(false);;
                    logger_1.default.error(`❌ Job failed`);
                    resolve({ success: false, error: 'Job failed' });
                }
                else {
                    // logger.info(`Percent: ${percent}% `);
                    // setProgressPercent(percent);
                }
            }, 1000);
        });
    }
    catch (error) {
        console.error('❌ Error running security test:', error);
    }
};
async function doTrusysSecurityTest(cmdObj) {
    (0, util_1.setupEnv)(cmdObj.envPath);
    const apiKey = cloud_1.cloudConfig.getApiKey() || (0, envars_1.getEnvString)('TRUSYS_API_KEY');
    const appUrl = cloud_1.cloudConfig.getTrusysAppUrl();
    const projectId = (0, accounts_1.getProjectId)();
    const workSpaceId = (0, accounts_1.getWorkSpaceId)();
    if (!projectId || !apiKey) {
        await cloud_1.cloudConfig.delete();
        logger_1.default.info(`❌ Not logged in. Run ${chalk_1.default.bold('trusys auth login')} to login.`);
        process.exitCode = 1;
        return null;
    }
    const apiHost = cloud_1.cloudConfig.getTrusysApiHost();
    // const apiHost = 'http://localhost:15500';
    const securityConfigName = cmdObj?.securityConfig;
    const securityConfigDetails = await getSecurityConfigDetails(securityConfigName, apiKey, apiHost);
    logger_1.default.info(`SecurityConfig: ${securityConfigDetails?.name}`);
    const app = await (0, eval_1.getApplicationDetils)(securityConfigDetails?.config?.target?.label, apiKey, apiHost);
    if (lodash_1.default.isEmpty(app)) {
        logger_1.default.error(`❌ Application not found`);
        return null;
    }
    const testSuite = {
        config: (0, sharedFrontend_1.getUnifiedConfig)(securityConfigDetails?.config, cmdObj?.applicationPath),
        force: true,
        verbose: false,
        delay: cmdObj?.delay ?? 0,
        maxConcurrency: cmdObj?.maxConcurrency ?? 20,
        applicationIds: [app?.id],
        securityConfigId: securityConfigDetails?.id,
        user: { email: cmdObj?.user },
        isCustomProvider: false,
    };
    if (app?.provider_config?.config?.custom_provider_type === 'Custom AI Application') {
        if (lodash_1.default.isEmpty(cmdObj?.applicationPath)) {
            logger_1.default.error(`❌ Please provide an application filepath for '${app?.name}' application`);
            return null;
        }
        else {
            const language = app?.provider_config?.config?.provider?.language;
            const applicationPath = cmdObj?.applicationPath;
            let hasCorrectExt = false;
            if (language === 'go') {
                hasCorrectExt = applicationPath?.endsWith('.go');
            }
            else if (language === 'python') {
                hasCorrectExt = applicationPath?.endsWith('.py');
            }
            else if (language === 'js') {
                hasCorrectExt = applicationPath?.endsWith('.js') || applicationPath?.endsWith('.ts');
            }
            if (!hasCorrectExt) {
                logger_1.default.error(`❌ Please give ${language} file path`);
                return null;
            }
            let configPath = cmdObj.config ?? 'trusysconfig.yaml';
            // If output filepath is not provided, locate the out file in the same directory as the config file:
            let redteamPath;
            if (cmdObj.output) {
                redteamPath = cmdObj.output;
            }
            else {
                const configDir = path.dirname(configPath);
                redteamPath = path.join(configDir, 'redteam.yaml');
            }
            if (testSuite?.config) {
                // Write liveRedteamConfig to a temporary file
                const filename = `redteam-${Date.now()}.yaml`;
                const tmpDir = cmdObj.loadedFromCloud ? '' : os.tmpdir();
                const tmpFile = path.join(tmpDir, filename);
                fs.mkdirSync(path.dirname(tmpFile), { recursive: true });
                fs.writeFileSync(tmpFile, js_yaml_1.default.dump(testSuite?.config));
                redteamPath = tmpFile;
                // Do not use default config.
                configPath = tmpFile;
                // logger.debug(`Using live config from ${tmpFile}`);
                // logger.debug(`Live config: ${JSON.stringify(cmdObj.liveRedteamConfig, null, 2)}`);
            }
            const redteamConfig = await (0, generate_1.doGenerateRedteam)({
                ...cmdObj,
                ...(cmdObj.liveRedteamConfig?.commandLineOptions || {}),
                config: configPath,
                output: redteamPath,
                inRedteamRun: true,
                abortSignal: cmdObj.abortSignal,
                progressBar: cmdObj.progressBar,
            });
            const providerId = typeof testSuite?.config?.targets[0]?.id === 'function'
                ? testSuite?.config.targets[0].id()
                : (testSuite?.config?.targets[0]?.id ?? '');
            const providerClass = await (0, providers_1.loadApiProvider)(providerId, {
                basePath: testSuite?.config?.targets[0]?.config?.basePath,
            });
            await Promise.all((redteamConfig?.tests ?? [])?.map(async (test) => {
                const prompt = {
                    raw: Array.isArray(redteamConfig?.prompts) ? redteamConfig?.prompts[0] : '',
                    label: Array.isArray(redteamConfig?.prompts)
                        ? redteamConfig?.prompts[0]
                        : '',
                };
                const renderedPrompt = await (0, evaluatorHelpers_1.renderPrompt)(prompt, test?.vars);
                const providerResponse = await providerClass.callApi(renderedPrompt);
                test.providerOutput = providerResponse.output;
            }));
            testSuite.config = redteamConfig;
        }
        testSuite.isCustomProvider = true;
    }
    const result = await securityTestRunHandler(testSuite, apiHost, apiKey);
    if (!result?.success) {
        logger_1.default.error(`Test suite failed: ${result?.error}`);
        return null;
    }
    logger_1.default.info(chalk_1.default.green(`Security run completed: ${appUrl}/workspaces/${workSpaceId}/projects/${projectId}/redteam-test-result/${result?.testRunId}`));
    return null;
}
function redteamRunCommand(program) {
    program
        .command('run')
        .description((0, dedent_1.default) `
        ${chalk_1.default.red('Red team')} a target application, a two-step process:

        1. Generates dynamic attack probes (i.e. test cases) tailored to your target application using specialized uncensored models.
        2. Evaluates the generated probes against your target application.
      `)
        .option('-c, --config [path]', 'Path to configuration file or cloud config UUID. Defaults to trusysconfig.yaml')
        .option('-o, --output [path]', 'Path to output file for generated tests. Defaults to redteam.yaml in the same directory as the configuration file.')
        .option('--no-cache', 'Do not read or write results to disk cache', false)
        .option('-j, --max-concurrency <number>', 'Maximum number of concurrent API calls', (val) => Number.parseInt(val, 10), evaluator_1.DEFAULT_MAX_CONCURRENCY)
        .option('--delay <number>', 'Delay in milliseconds between API calls', (val) => Number.parseInt(val, 10))
        .option('--remote', 'Force remote inference wherever possible', false)
        .option('--force', 'Force generation even if no changes are detected', false)
        .option('--no-progress-bar', 'Do not show progress bar')
        .option('--filter-providers, --filter-targets <providers>', 'Only run tests with these providers (regex match)')
        .option('--securityConfig <securityConfig>', 'Security Config to use')
        .option('--user <user>', 'Email Id for Security Run')
        .option('-t, --target <id>', 'Cloud provider target ID to run the scan on')
        .option('--applicationPath <path>', 'Path to custom provider file')
        .action(async (opts) => {
        (0, util_1.setupEnv)(opts.envPath);
        telemetry_1.default.record('command_used', {
            name: 'redteam run',
        });
        await telemetry_1.default.send();
        if (opts.config && UUID_REGEX.test(opts.config)) {
            if (opts.target && !UUID_REGEX.test(opts.target)) {
                throw new Error('Invalid target ID, it must be a valid UUID');
            }
            const configObj = await (0, cloud_2.getConfigFromCloud)(opts.config, opts.target);
            // backwards compatible for old cloud servers
            if (opts.target &&
                UUID_REGEX.test(opts.target) &&
                (!configObj.targets || configObj.targets?.length === 0)) {
                configObj.targets = [{ id: `${constants_1.CLOUD_PROVIDER_PREFIX}${opts.target}`, config: {} }];
            }
            opts.liveRedteamConfig = configObj;
            opts.config = undefined;
            opts.loadedFromCloud = true;
        }
        else if (opts.target) {
            logger_1.default.error(`Target ID (-t) can only be used when -c is used. To use a cloud target inside of a config set the id of the target to ${constants_1.CLOUD_PROVIDER_PREFIX}${opts.target}. `);
            process.exitCode = 1;
            return;
        }
        try {
            if (opts.securityConfig) {
                if (opts.user) {
                    const EMAIL_REGEX = /^(([^<>()[\]\\.,;:\s@"]+(\.[^<>()[\]\\.,;:\s@"]+)*)|(".+"))@((\[[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\])|(([a-zA-Z\-0-9]+\.)+[a-zA-Z]{2,}))$/;
                    if (!EMAIL_REGEX.test(opts.user)) {
                        logger_1.default.error('Please provide a valid email id using --user <user> to run security.');
                        return;
                    }
                    await doTrusysSecurityTest(opts);
                }
                else {
                    logger_1.default.error('Please provide an email id to run security.');
                }
                return;
            }
            if (opts.remote) {
                cliState_1.default.remote = true;
            }
            await (0, shared_1.doRedteamRun)(opts);
        }
        catch (error) {
            if (error instanceof zod_1.z.ZodError) {
                logger_1.default.error('Invalid options:');
                error.errors.forEach((err) => {
                    logger_1.default.error(`  ${err.path.join('.')}: ${err.message}`);
                });
            }
            else {
                logger_1.default.error(`An unexpected error occurred during red team run: ${error instanceof Error ? error.message : String(error)}\n${error instanceof Error ? error.stack : ''}`);
            }
            process.exitCode = 1;
        }
    });
    (0, poison_1.poisonCommand)(program);
}
//# sourceMappingURL=run.js.map