gfwleak/enderbyendera-realtime-protection
Archived
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

1、任务 添加 附件上传功能

Browse Source
This commit is contained in:
PushM
2024-07-18 10:32:19 +08:00
parent d9605328ca
commit 9b53d78fb6
7 changed files with 242 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/main/java/com/realtime/protection/configuration/utils/File.java
View File

@@ -12,4 +12,5 @@ public class File implements Serializable {
private Long fileSize;
private String fileType;
private Long staticRuleId;
private Long taskId;
}

8
src/main/java/com/realtime/protection/server/rule/staticrule/StaticRuleController.java
View File

@@ -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();

83
src/main/java/com/realtime/protection/server/task/TaskController.java
View File

@@ -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<File> 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<Resource> 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,

101
src/main/java/com/realtime/protection/server/task/TaskControllerApi.java
View File

@@ -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<Resource> downFile(@RequestParam Integer id) throws MalformedURLException, UnsupportedEncodingException;
@GetMapping("/query")
@Operation(
summary = "查询任务",
description = "按页和搜索内容查询任务相关信息",

7
src/main/java/com/realtime/protection/server/task/TaskMapper.java
View File

@@ -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<Task> queryAuditInfoNotification(String userId , Integer auditStatus);
void updateNotificationByTaskId(@Param("taskIds") List<Long> taskIdList);
void saveFile(File saveFile);
File selectFileById(Integer id);
List<File> selectFilesByStaticRuleId(Long taskId);
}

15
src/main/java/com/realtime/protection/server/task/TaskService.java
View File

@@ -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<File> selectFilesByStaticRuleId(Long taskId) {
return taskMapper.selectFilesByStaticRuleId(taskId);
}
public void saveFile(File saveFile) {
taskMapper.saveFile(saveFile);
}
public File selectFileById(Integer id) {
return taskMapper.selectFileById(id);
}
}