diff --git a/src/main/java/com/realtime/protection/configuration/utils/File.java b/src/main/java/com/realtime/protection/configuration/utils/File.java index 38a8e8c..c54db65 100644 --- a/src/main/java/com/realtime/protection/configuration/utils/File.java +++ b/src/main/java/com/realtime/protection/configuration/utils/File.java @@ -12,4 +12,5 @@ public class File implements Serializable { private Long fileSize; private String fileType; private Long staticRuleId; + private Long taskId; } diff --git a/src/main/java/com/realtime/protection/server/rule/staticrule/StaticRuleController.java b/src/main/java/com/realtime/protection/server/rule/staticrule/StaticRuleController.java index 24c1ce5..cf3c002 100644 --- a/src/main/java/com/realtime/protection/server/rule/staticrule/StaticRuleController.java +++ b/src/main/java/com/realtime/protection/server/rule/staticrule/StaticRuleController.java @@ -89,7 +89,7 @@ public class StaticRuleController implements StaticRuleControllerApi { public ResponseResult uploadAttachment(@RequestParam("static_rule_id") Long staticRuleId, MultipartFile file) throws IOException { - String uploadPath = "d:\\"; + String uploadPath = "/static_rule_file"; //获取文件上传名称 String fileName=file.getOriginalFilename(); //获取文件保存全路径 @@ -99,6 +99,12 @@ public class StaticRuleController implements StaticRuleControllerApi { //获取文件类型 String fileType=file.getContentType(); java.io.File newFile=new java.io.File(savePath); + // 获取文件的父目录 + java.io.File parentDir = newFile.getParentFile(); + // 如果父目录不存在，则创建该目录 + if (parentDir != null && !parentDir.exists()) { + parentDir.mkdirs(); + } //TODO 注意要将文件保存到本地路径中 file.transferTo(newFile); File saveFile=new File(); diff --git a/src/main/java/com/realtime/protection/server/task/TaskController.java b/src/main/java/com/realtime/protection/server/task/TaskController.java index 55351cb..51f6025 100644 --- a/src/main/java/com/realtime/protection/server/task/TaskController.java +++ b/src/main/java/com/realtime/protection/server/task/TaskController.java @@ -5,6 +5,7 @@ import com.realtime.protection.configuration.entity.task.TaskCommandInfo; import com.realtime.protection.configuration.entity.user.UserFull; import com.realtime.protection.configuration.exception.DorisStartException; import com.realtime.protection.configuration.response.ResponseResult; +import com.realtime.protection.configuration.utils.File; import com.realtime.protection.configuration.utils.enums.StateEnum; import com.realtime.protection.configuration.utils.enums.audit.AuditStatusEnum; import com.realtime.protection.server.command.CommandService; @@ -21,10 +22,21 @@ import jakarta.validation.constraints.Max; import jakarta.validation.constraints.Min; import jakarta.validation.constraints.NotNull; import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.core.io.Resource; +import org.springframework.core.io.UrlResource; +import org.springframework.http.HttpHeaders; import org.springframework.http.HttpStatus; +import org.springframework.http.MediaType; import org.springframework.http.ResponseEntity; import org.springframework.web.bind.annotation.*; import org.springframework.web.multipart.MultipartFile; + +import java.io.IOException; +import java.io.UnsupportedEncodingException; +import java.net.MalformedURLException; +import java.net.URI; +import java.net.URLEncoder; +import java.nio.file.Paths; import java.util.Base64; @@ -139,6 +151,77 @@ public class TaskController implements TaskControllerApi { .setData("command_hash",taskCommandInfo.hashCode()); } + @GetMapping("/queryAttachment") + @Override + public ResponseResult queryAttachment(@RequestParam("task_id") Long taskId) { + List files = taskService.selectFilesByStaticRuleId(taskId); + + return ResponseResult.ok() + .setData("files", files); + } + + //上传附件 + @PostMapping("/uploadAttachment") + @Override + public ResponseResult uploadAttachment(@RequestParam("task_id") Long taskId, + MultipartFile file) throws IOException { + +// String uploadPath = "d:\\"; + String uploadPath = "/task_file"; + //获取文件上传名称 + String fileName=file.getOriginalFilename(); + //获取文件保存全路径 + String savePath=uploadPath+"/"+fileName; + //获取文件大小 + Long fileSize=file.getSize(); + //获取文件类型 + String fileType=file.getContentType(); + java.io.File newFile=new java.io.File(savePath); + // 获取文件的父目录 + java.io.File parentDir = newFile.getParentFile(); + // 如果父目录不存在，则创建该目录 + if (parentDir != null && !parentDir.exists()) { + parentDir.mkdirs(); + } + //TODO 注意要将文件保存到本地路径中 + file.transferTo(newFile); + File saveFile=new File(); + saveFile.setFileName(fileName); + saveFile.setFileSize(fileSize); + saveFile.setFileType(fileType); + saveFile.setFilePath(savePath); + saveFile.setTaskId(taskId); + taskService.saveFile(saveFile); + + return ResponseResult.ok(); + } + + + + //文件下载 + @GetMapping("downAttachment") + @Override + public ResponseEntity downFile(@RequestParam Integer id) throws MalformedURLException, UnsupportedEncodingException { + //首先根据id，从数据库获取文件信息 + File downFile=taskService.selectFileById(id); + if (downFile!=null){ + String path= downFile.getFilePath(); + //本地路径地址转为url编码路径 + URI urlPath= Paths.get(path).toUri(); + Resource resource=new UrlResource(urlPath); + if (resource.exists()){ + HttpHeaders headers=new HttpHeaders(); + headers.add(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_OCTET_STREAM+""); + headers.add(HttpHeaders.CONTENT_LENGTH,downFile.getFileSize()+""); + //注意文件名处要改为URL编码 + headers.add(HttpHeaders.CONTENT_DISPOSITION,"attachment; filename=\"" + + URLEncoder.encode(downFile.getFileName(), "utf-8") + "\""); + return ResponseEntity.ok().headers(headers).body(resource); + } + } + return ResponseEntity.notFound().build(); + } + @Override @GetMapping("/query") public ResponseResult queryTasks(@RequestParam(value = "task_status", required = false) Integer taskStatus, diff --git a/src/main/java/com/realtime/protection/server/task/TaskControllerApi.java b/src/main/java/com/realtime/protection/server/task/TaskControllerApi.java index 3c10190..29cd3f7 100644 --- a/src/main/java/com/realtime/protection/server/task/TaskControllerApi.java +++ b/src/main/java/com/realtime/protection/server/task/TaskControllerApi.java @@ -17,10 +17,14 @@ import jakarta.validation.constraints.Max; import jakarta.validation.constraints.Min; import jakarta.validation.constraints.NotNull; import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.core.io.Resource; import org.springframework.http.ResponseEntity; import org.springframework.web.bind.annotation.*; import org.springframework.web.multipart.MultipartFile; +import java.io.IOException; +import java.io.UnsupportedEncodingException; +import java.net.MalformedURLException; import java.security.InvalidKeyException; import java.security.NoSuchAlgorithmException; import java.time.LocalDate; @@ -107,8 +111,103 @@ public interface TaskControllerApi { ResponseResult newTaskWithAPI(@RequestBody @Valid TaskCommandInfo taskCommandInfo, @Autowired HttpServletRequest request) throws NoSuchAlgorithmException, InvalidKeyException ; + @Operation( + summary = "查询任务的附件信息", + description = "查询任务的所有附件信息", + responses = { + @io.swagger.v3.oas.annotations.responses.ApiResponse( + description = "返回任务的所有附件信息", + content = @Content( + mediaType = "application/json", + schema = @Schema( + implementation = ResponseResult.class), + examples = @ExampleObject( + name = "查询任务的附件信息", + value = """ + { + { + "code": 200, + "message": "请求成功", + "data": { + "files": [ + { + "id": 3, + "fileName": "系统角色.docx", + "filePath": null, + "fileSize": 10915, + "fileType": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "taskId": 1819 + }, + { + "id": 4, + "fileName": "权限.txt", + "filePath": null, + "fileSize": 100, + "fileType": "text/plain", + "taskId": 1819 + } + ] + }, + "another": null + } + }""", + description = "task_id:任务id" + + "id:文件id" + "fileName:文件名" + "filePath:文件路径" + + "fileSize:文件大小" + "fileType:文件类型" + "taskId:taskIdid" + ) + ) + ) + }, + parameters = { + @Parameter(name = "task_id", description = "任务id") + } + ) + @GetMapping("/queryAttachment") + ResponseResult queryAttachment(@RequestParam("task_id") Long taskId); - @GetMapping("/query") + @Operation( + summary = "上传附件", + description = "以文件方式上传附件", + responses = { + @io.swagger.v3.oas.annotations.responses.ApiResponse( + description = "返回上传结果", + content = @Content( + mediaType = "application/json", + schema = @Schema(implementation = ResponseResult.class) + ) + ) + }, + parameters = { + @Parameter(name = "task_id", description = "任务id"), + } + + ) + //上传附件 + @PostMapping("/uploadAttachment") + ResponseResult uploadAttachment(@RequestParam("task_id") Long taskId, + MultipartFile file) throws IOException; + + //文件下载 + @Operation( + summary = "下载附件", + description = "下载附件", + responses = { + @io.swagger.v3.oas.annotations.responses.ApiResponse( + description = "返回下载结果", + content = @Content( + mediaType = "application/octet-stream", + schema = @Schema(implementation = ResponseEntity.class) + ) + ) + }, + parameters = { + @Parameter(name = "id", description = "文件id") + } + ) + @GetMapping("downAttachment") + ResponseEntity downFile(@RequestParam Integer id) throws MalformedURLException, UnsupportedEncodingException; + + @GetMapping("/query") @Operation( summary = "查询任务", description = "按页和搜索内容查询任务相关信息", diff --git a/src/main/java/com/realtime/protection/server/task/TaskMapper.java b/src/main/java/com/realtime/protection/server/task/TaskMapper.java index a0abe01..8653551 100644 --- a/src/main/java/com/realtime/protection/server/task/TaskMapper.java +++ b/src/main/java/com/realtime/protection/server/task/TaskMapper.java @@ -4,6 +4,7 @@ import com.realtime.protection.configuration.entity.defense.object.ProtectObject import com.realtime.protection.configuration.entity.task.DynamicTaskInfo; import com.realtime.protection.configuration.entity.task.Task; import com.realtime.protection.configuration.entity.task.TaskCommandInfo; +import com.realtime.protection.configuration.utils.File; import org.apache.ibatis.annotations.Mapper; import org.apache.ibatis.annotations.Param; import org.apache.ibatis.annotations.Select; @@ -132,4 +133,10 @@ public interface TaskMapper { List queryAuditInfoNotification(String userId , Integer auditStatus); void updateNotificationByTaskId(@Param("taskIds") List taskIdList); + + void saveFile(File saveFile); + + File selectFileById(Integer id); + + List selectFilesByStaticRuleId(Long taskId); } diff --git a/src/main/java/com/realtime/protection/server/task/TaskService.java b/src/main/java/com/realtime/protection/server/task/TaskService.java index ef7cfd2..4d69297 100644 --- a/src/main/java/com/realtime/protection/server/task/TaskService.java +++ b/src/main/java/com/realtime/protection/server/task/TaskService.java @@ -9,6 +9,7 @@ import com.realtime.protection.configuration.entity.task.DynamicTaskInfo; import com.realtime.protection.configuration.entity.task.Task; import com.realtime.protection.configuration.entity.task.TaskCommandInfo; import com.realtime.protection.configuration.utils.Counter; +import com.realtime.protection.configuration.utils.File; import com.realtime.protection.configuration.utils.SqlSessionWrapper; import com.realtime.protection.configuration.utils.enums.StateEnum; import com.realtime.protection.configuration.utils.enums.TaskTypeEnum; @@ -646,6 +647,18 @@ public class TaskService { "事件类型：APT攻击 ，开始时间"+task.getTaskStartTime()+"结束时间："+task.getTaskEndTime()); } - } + + public List selectFilesByStaticRuleId(Long taskId) { + return taskMapper.selectFilesByStaticRuleId(taskId); + } + public void saveFile(File saveFile) { + taskMapper.saveFile(saveFile); + } + + public File selectFileById(Integer id) { + return taskMapper.selectFileById(id); + } + + } diff --git a/src/main/resources/mappers/TaskMapper.xml b/src/main/resources/mappers/TaskMapper.xml index ac36640..51ddb7c 100644 --- a/src/main/resources/mappers/TaskMapper.xml +++ b/src/main/resources/mappers/TaskMapper.xml @@ -492,6 +492,14 @@ #{taskId} + + insert into t_task_file( + file_name, file_path, file_size, file_type, task_id + ) + values ( + #{fileName}, #{filePath}, #{fileSize}, #{fileType}, #{taskId} + ) + + + + + + + + + + + + + + + \ No newline at end of file