diff --git a/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/engine/worker/afl/AFLSlotCompileWorker.java b/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/engine/worker/afl/AFLSlotCompileWorker.java index 6c2256bf..3b9247c4 100644 --- a/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/engine/worker/afl/AFLSlotCompileWorker.java +++ b/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/engine/worker/afl/AFLSlotCompileWorker.java @@ -129,10 +129,10 @@ public class AFLSlotCompileWorker extends DockerWorker { "cd "+prefix+"\n"; return cmd; } - if ("gz".equals(suffix)) { + if (fileName.contains(".tar")||fileName.contains(".tar.gz")) { String cmd ="tar -zxvf "+fileName+"\n" + "chmod -R o+rwx "+prefix+"/\n"+ - "cd "+prefix+"\n"; + "cd "+fileName.replace(".tar.gz","").replace(".tar","")+"\n"; return cmd; } return ""; diff --git a/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/engine/worker/afl/AFLWorker.java b/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/engine/worker/afl/AFLWorker.java index fac19984..72b06a8f 100644 --- a/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/engine/worker/afl/AFLWorker.java +++ b/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/engine/worker/afl/AFLWorker.java @@ -115,9 +115,12 @@ public class AFLWorker extends DockerWorker { } } private String cdSourceName(String fileName){ - int dotIndex = fileName.lastIndexOf("."); - if (dotIndex!=-1) { - return "cd "+fileName.substring(0, dotIndex)+"\n"; + if (fileName.contains(".zip")||fileName.contains(".tar")||fileName.contains(".tar.gz")) { + String cdPath = fileName + .replace(".zip","") + .replace(".tar.gz","") + .replace(".tar",""); + return "cd "+cdPath+"\n"; } return ""; } diff --git a/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/engine/worker/afl/TestCaseGenerationWorker.java b/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/engine/worker/afl/TestCaseGenerationWorker.java index be4935ff..c599563a 100644 --- a/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/engine/worker/afl/TestCaseGenerationWorker.java +++ b/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/engine/worker/afl/TestCaseGenerationWorker.java @@ -2,7 +2,7 @@ package cd.casic.ci.process.engine.worker.afl; import cd.casic.ci.process.common.WorkAtom; -import cd.casic.ci.process.dto.req.resource.ResourceQueryReq; +import cd.casic.ci.process.dto.req.testCase.TestCaseAITaskCreateReq; import cd.casic.ci.process.dto.resp.resource.ResourceDetailResp; import cd.casic.ci.process.engine.runContext.TaskRunContext; import cd.casic.ci.process.engine.worker.base.DockerWorker; @@ -11,20 +11,24 @@ import cd.casic.ci.process.process.dataObject.target.TargetVersion; import cd.casic.ci.process.process.dataObject.task.PipTask; import cd.casic.ci.process.process.dataObject.testCase.TestCaseInfo; import cd.casic.ci.process.process.service.target.TargetVersionService; +import cd.casic.ci.process.process.service.testCase.TestCaseAIGeneratorService; import cd.casic.ci.process.process.service.testCase.TestCaseInfoService; import cd.casic.ci.process.util.SftpUploadUtil; import cd.casic.module.execute.docker.dataobject.dto.DockerEndpointDo; import cd.casic.module.machine.dal.dataobject.MachineInfoDO; import cd.casic.module.machine.utils.CryptogramUtil; import com.alibaba.fastjson.JSON; +import com.alibaba.fastjson.JSONObject; +import com.jcraft.jsch.*; import jakarta.annotation.Resource; import lombok.extern.slf4j.Slf4j; import org.apache.commons.lang3.StringUtils; -import java.io.File; +import java.io.*; import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Properties; import static cd.casic.ci.process.engine.constant.PipelineVariableConstant.*; import static cd.casic.ci.process.engine.constant.TestCaseGenerationConstant.*; @@ -36,6 +40,8 @@ public class TestCaseGenerationWorker extends DockerWorker { private TestCaseInfoService testCaseInfoService; @Resource private TargetVersionService targetVersionService; + @Resource + private TestCaseAIGeneratorService aiGeneratorService; @Override public void execute(TaskRunContext context) { int statusCode = -1; @@ -45,7 +51,6 @@ public class TestCaseGenerationWorker extends DockerWorker { String seedPath= workDir+File.separator+ AFL_DOCKER_SEED; Map taskProperties = taskDef.getTaskProperties(); Object caseType = taskProperties.get(CASE_TYPE_KEY); - String binaryName = taskProperties.get(BINARY_NAME) instanceof String ? ((String) taskProperties.get("binaryName")) : null; String managerId = getVariableNearby(context, AFL_RESOURCE_MANAGER_ID_KEY) instanceof String ? ((String) getVariableNearby(context, AFL_RESOURCE_MANAGER_ID_KEY)) : null; Object itemListObject = taskProperties.get(ITEM_LIST); List itemList = JSON.parseArray(JSON.toJSONString(itemListObject),String.class); @@ -55,23 +60,21 @@ public class TestCaseGenerationWorker extends DockerWorker { if (Objects.isNull(machineInfo)||Objects.isNull(dockerInfo)) { append(context,"该资源不支持docker或者ssh"); } - PipPipeline pipeline = (PipPipeline) getContextManager().getContext(taskDef.getPipelineId()).getContextDef(); - TargetVersion targetVersion = targetVersionService.getById(pipeline.getTargetVersionId()); - String fileName = targetVersion.getFileName(); - String imageName = getVariableNearby(context,IMAGE_NAME) instanceof String ? ((String) getVariableNearby(context,IMAGE_NAME)) : null; + //如果machineId为0，则说明该节点没有配置机器，则使用开始节点的机器 if (CASE_TYPE_AI.equals(caseType)) { - // 容器内部test 目录, 获取编译后的文件 - String runScript = "docker run -v "+ AFL_VOLUME_WORK_DIR_PREFIX +":"+ AFL_DOCKER_BASE_DIR+" -it "+imageName+" bash\n" // 目前测试用例ai生成只有这个路径能跑 - + "cd "+ AFL_DOCKER_BASE_DIR+"\n" + - "mkdir -p "+seedPath +"\n"+ - "cd "+seedPath +"\n"+ - "PYTHONPATH="+AFL_DOCKER_BASE_DIR+File.separator+"CaseGenerator/src python3 "+AFL_DOCKER_BASE_DIR+File.separator+"CaseGenerator/src/CaseGenerator/main.py --work-dir "+workDir+" --binary "+workDir+File.separator+getSourceName(fileName)+File.separator+binaryName+" --output-dir "+seedPath+" --count 100"; - //将节点的配置信息反编译成对象 - log.info("测试用例生成-AI生成，实际执行命令:{}" , runScript); - append(context,"测试用例生成-AI生成，实际执行命令:"+runScript); - //获取机器 - dockerRun(runScript,dockerInfo,context); +// // 容器内部test 目录, 获取编译后的文件 +// String runScript = "docker run -v "+ AFL_VOLUME_WORK_DIR_PREFIX +":"+ AFL_DOCKER_BASE_DIR+" -it "+imageName+" bash\n" // 目前测试用例ai生成只有这个路径能跑 +// + "cd "+ AFL_DOCKER_BASE_DIR+"\n" + +// "mkdir -p "+seedPath +"\n"+ +// "cd "+seedPath +"\n"+ +// "PYTHONPATH="+AFL_DOCKER_BASE_DIR+File.separator+"CaseGenerator/src python3 "+AFL_DOCKER_BASE_DIR+File.separator+"CaseGenerator/src/CaseGenerator/main.py --work-dir "+workDir+" --binary "+workDir+File.separator+getSourceName(fileName)+File.separator+binaryName+" --output-dir "+seedPath+" --count 100"; +// //将节点的配置信息反编译成对象 +// log.info("测试用例生成-AI生成，实际执行命令:{}" , runScript); +// append(context,"测试用例生成-AI生成，实际执行命令:"+runScript); +// //获取机器 +// dockerRun(runScript,dockerInfo,context); + handleAIGenerator(context); } else if (CASE_TYPE_FILE.equals(caseType)){ // 文件上传 Object filePathObj = taskProperties.get(SEED_SOURCE); @@ -121,10 +124,89 @@ public class TestCaseGenerationWorker extends DockerWorker { globalVariables.put(AFL_DOCKER_SEED_PATH_KEY,seedPath); } } - public String getSourceName(String fileName){ - int dotIndex = fileName.lastIndexOf("."); - if (dotIndex!=-1) { - return fileName.substring(0, dotIndex); + public void handleAIGenerator(TaskRunContext context){ + PipTask taskDef = context.getContextDef() instanceof PipTask ? ((PipTask) context.getContextDef()) : null; + Map taskProperties = taskDef.getTaskProperties(); + String binaryName = taskProperties.get(BINARY_NAME) instanceof String ? ((String) taskProperties.get("binaryName")) : null; + PipPipeline pipeline = (PipPipeline) getContextManager().getContext(taskDef.getPipelineId()).getContextDef(); + TargetVersion targetVersion = targetVersionService.getById(pipeline.getTargetVersionId()); + String managerId = getVariableNearby(context, AFL_RESOURCE_MANAGER_ID_KEY) instanceof String ? ((String) getVariableNearby(context, AFL_RESOURCE_MANAGER_ID_KEY)) : null; + ResourceDetailResp resourceById = resourceManagerService.findResourceDetailById(managerId); + MachineInfoDO machineInfo = resourceById.getMachineInfo(); + DockerEndpointDo dockerInfo = resourceById.getDockerInfo(); + String fileName = targetVersion.getFileName(); + if (Objects.isNull(machineInfo)||Objects.isNull(dockerInfo)) { + append(context,"该资源不支持docker或者ssh"); + } + // TODO 需要新增的项 + String prompt = "你是一个模糊测试专家，当前程序输入为文本例如 aa aaa 这种"; + Integer count = 10; + // 获取编译后的文文件的inputStream + JSch jsch = new JSch(); + Session session = null; + try { + session = jsch.getSession(machineInfo.getUsername(),machineInfo.getHostIp(), machineInfo.getSshPort()); + session.setPassword( CryptogramUtil.doDecrypt(machineInfo.getPassword())); + // 跳过 host key 检查 + Properties config = new Properties(); + config.put("StrictHostKeyChecking", "no"); + session.setConfig(config); + session.connect(); + ChannelSftp sftp = (ChannelSftp) session.openChannel("sftp"); + sftp.connect(); + String workDir = (String)getGlobalVariable(context,AFL_DOCKER_WORK_DIR_KEY); + String filePath = AFL_DOCKER_BASE_DIR + File.separator + workDir+File.separator+getSourceName(fileName)+File.separator+binaryName; + InputStream inputStream = sftp.get(filePath); + String targetUploadRes = aiGeneratorService.targetUpload("./" + binaryName, inputStream); + String binaryId = JSON.parseObject(targetUploadRes).getString("id"); + // 创建task + TestCaseAITaskCreateReq req = new TestCaseAITaskCreateReq(); + req.setBinaryId(binaryId); + req.setCustomPrompt(prompt); + req.setCount(count); + String taskCreateRes = aiGeneratorService.taskCreate(req); + String taskId = JSON.parseObject(taskCreateRes).getString("task_id"); + while (true){ + // 获取状态及日志 + String taskStatusRes = aiGeneratorService.getTaskStatus(taskId); + String status = JSONObject.parseObject(taskStatusRes).getString("status"); + if ("success".equals(status)) { + break; + } else if("failed".equals(status)){ + toBadEnding(); + } else { + String logRes = aiGeneratorService.getTaskLog(taskId); + String logStr = JSON.parseObject(logRes).getString("log"); + append(context,logStr); + } + } + File tempFile = File.createTempFile("seed", ".zip"); + tempFile.deleteOnExit(); + FileOutputStream fos = new FileOutputStream(tempFile); + // 下载生成的种子 + aiGeneratorService.downloadTestCase(taskId,fos); + fos.close(); + FileInputStream fis = new FileInputStream(tempFile); + // 上传种子zip到目标服务器 + sftp.put(fis,""); + sftp.disconnect(); + String seedPath= workDir+File.separator+ AFL_DOCKER_SEED; + String seedTarget = AFL_VOLUME_WORK_DIR_PREFIX+File.separator+seedPath+File.separator+"seed.zip"; + String seedDir = AFL_VOLUME_WORK_DIR_PREFIX+File.separator+seedPath; + // 解压种子zip + ChannelExec exec = (ChannelExec)session.openChannel("exec"); + exec.setCommand("unzip -o "+seedTarget+" -d "+seedDir); + } catch (Exception e) { + append(context,"测试用例生成失败"); + } + + } + private String getSourceName(String fileName){ + if (fileName.contains(".zip")||fileName.contains(".tar")||fileName.contains(".tar.gz")) { + return fileName + .replace(".zip","") + .replace(".tar.gz","") + .replace(".tar",""); } return ""; } diff --git a/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/process/service/testCase/impl/TestCaseAIGeneratorServiceImpl.java b/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/process/service/testCase/impl/TestCaseAIGeneratorServiceImpl.java index c501fad8..dfea6ddc 100644 --- a/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/process/service/testCase/impl/TestCaseAIGeneratorServiceImpl.java +++ b/modules/module-ci-process-biz/src/main/java/cd/casic/ci/process/process/service/testCase/impl/TestCaseAIGeneratorServiceImpl.java @@ -43,6 +43,7 @@ public class TestCaseAIGeneratorServiceImpl implements TestCaseAIGeneratorServic zos.putNextEntry(entry); IOUtils.copy(inputStream,zos); zos.closeEntry(); + zos.close(); inputStream.close(); FileSystemResource fsr = new FileSystemResource(tempFile); String url = properties.getBaseUrl()+ TestCaseAIGeneratorConstant.targetUpload; @@ -100,7 +101,8 @@ public class TestCaseAIGeneratorServiceImpl implements TestCaseAIGeneratorServic headers.setContentType(MediaType.APPLICATION_JSON); HttpEntity entity = new HttpEntity<>("{}",headers); String url = properties.getBaseUrl() + TestCaseAIGeneratorConstant.downloadTestCase+taskId; - ResponseEntity exchange = restTemplate.exchange(url, HttpMethod.GET, entity, InputStream.class, new HashMap<>()); - IOUtils.copy(exchange.getBody(),outputStream); + ResponseEntity exchange = restTemplate.exchange(url, HttpMethod.GET, entity, byte[].class, new HashMap<>()); + outputStream.write(exchange.getBody()); } + } diff --git a/ops-server/src/test/java/cd/casic/server/TestCaseAIGeneratorTest.java b/ops-server/src/test/java/cd/casic/server/TestCaseAIGeneratorTest.java new file mode 100644 index 00000000..eb12aef2 --- /dev/null +++ b/ops-server/src/test/java/cd/casic/server/TestCaseAIGeneratorTest.java @@ -0,0 +1,66 @@ +package cd.casic.server; + +import cd.casic.ci.process.constant.TestCaseAIGeneratorConstant; +import cd.casic.ci.process.dto.req.testCase.TestCaseAITaskCreateReq; +import cd.casic.ci.process.process.service.testCase.TestCaseAIGeneratorService; +import jakarta.annotation.Resource; +import lombok.extern.slf4j.Slf4j; +import org.junit.jupiter.api.Test; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.core.io.ClassPathResource; +import org.springframework.test.context.ActiveProfiles; +import org.springframework.test.context.web.WebAppConfiguration; + +import java.io.File; +import java.io.FileOutputStream; +import java.io.IOException; + + +@SpringBootTest(classes = {OpsServerApplication.class}) +@ActiveProfiles("local") +@Slf4j +@WebAppConfiguration +public class TestCaseAIGeneratorTest { + @Resource + private TestCaseAIGeneratorService service; + + @Test + public void targetUpload() throws IOException { + //{"success":true,"msg":"","id":"8711fd176b686eebcc4e3f49bdde1a18"} + ClassPathResource classPathResource = new ClassPathResource("static/pcre2test"); + String s = service.targetUpload("./pcre2test", classPathResource.getInputStream()); + System.out.println(s); + } + @Test + public void taskCreate(){ + // {"success":true,"msg":"","task_id":"7c5d7df0-68d5-4935-aae7-00287976ae44"} +// {"success":true,"msg":"","task_id":"a70f92e9-9b8c-4be3-974f-8b0efa253237"} + TestCaseAITaskCreateReq req = new TestCaseAITaskCreateReq(); + req.setBinaryId("8711fd176b686eebcc4e3f49bdde1a18"); + req.setCustomPrompt("你是一个模糊测试专家，当前程序输入为文本例如 aa aaa 这种"); + req.setCount(10); + String s = service.taskCreate(req); + System.out.println(s); + } + @Test + public void getTaskStatus(){ +// {"success":true,"msg":"","task_id":"7c5d7df0-68d5-4935-aae7-00287976ae44","status":"success"} + String taskStatus = service.getTaskStatus("1e6e9694-5ae2-425e-ad80-9af266d33bea"); + System.out.println(taskStatus); + } + @Test + public void getTaskLog(){ +// {"success":true,"msg":"","log":"智能体初始化中...\n"} + String taskLog = service.getTaskLog("1e6e9694-5ae2-425e-ad80-9af266d33bea"); + System.out.println(taskLog); + } + @Test + public void downloadSeed() throws IOException { + File tempFile = File.createTempFile("seed", ".zip"); + FileOutputStream fos = new FileOutputStream(tempFile); + service.downloadTestCase("1e6e9694-5ae2-425e-ad80-9af266d33bea",fos); + tempFile.deleteOnExit(); + fos.close(); + System.out.println(); + } +}