测试用例生成以及AFL逻辑修改

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 "";

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 (fileName.contains(".zip")||fileName.contains(".tar")||fileName.contains(".tar.gz")) {
-        if (dotIndex!=-1) {
+            String cdPath = fileName
-            return "cd "+fileName.substring(0, dotIndex)+"\n";
+                    .replace(".zip","")
+                    .replace(".tar.gz","")
+                    .replace(".tar","");
+            return "cd "+cdPath+"\n";
         }
         return "";
     }

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<String, Object> 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<String> 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 目录, 获取编译后的文件
+//                // 容器内部test 目录, 获取编译后的文件
-                String runScript = "docker run -v "+ AFL_VOLUME_WORK_DIR_PREFIX +":"+ AFL_DOCKER_BASE_DIR+"  -it "+imageName+" bash\n" // 目前测试用例ai生成只有这个路径能跑
+//                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" +
+//                + "cd "+ AFL_DOCKER_BASE_DIR+"\n" +
-                        "mkdir -p "+seedPath +"\n"+
+//                        "mkdir -p "+seedPath +"\n"+
-                        "cd "+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";
+//                        "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);
+//                    log.info("测试用例生成-AI生成，实际执行命令:{}" , runScript);
-                    append(context,"测试用例生成-AI生成，实际执行命令:"+runScript);
+//                    append(context,"测试用例生成-AI生成，实际执行命令:"+runScript);
-                    //获取机器
+//                    //获取机器
-                dockerRun(runScript,dockerInfo,context);
+//                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){
+    public void handleAIGenerator(TaskRunContext context){
-        int dotIndex = fileName.lastIndexOf(".");
+        PipTask taskDef = context.getContextDef() instanceof PipTask ? ((PipTask) context.getContextDef()) : null;
-        if (dotIndex!=-1) {
+        Map<String, Object> taskProperties = taskDef.getTaskProperties();
-            return fileName.substring(0, dotIndex);
+        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 "";
     }

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<String> entity = new HttpEntity<>("{}",headers);
         String url = properties.getBaseUrl() + TestCaseAIGeneratorConstant.downloadTestCase+taskId;
-        ResponseEntity<InputStream> exchange = restTemplate.exchange(url, HttpMethod.GET, entity, InputStream.class, new HashMap<>());
+        ResponseEntity<byte[]> exchange = restTemplate.exchange(url, HttpMethod.GET, entity, byte[].class, new HashMap<>());
-        IOUtils.copy(exchange.getBody(),outputStream);
+        outputStream.write(exchange.getBody());
     }
+
 }

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();
+    }
+}