feat: ASW-8 新增 Runner 相关接口

This commit is contained in:
shizhendong
2024-07-16 17:38:32 +08:00
parent 5a4c15b00a
commit 79146845b9
15 changed files with 478 additions and 3 deletions

View File

@@ -1,6 +1,7 @@
package net.geedge.asw.common.util;
import cn.hutool.core.date.DateTime;
import cn.hutool.log.Log;
import com.baomidou.mybatisplus.core.metadata.OrderItem;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
@@ -11,10 +12,13 @@ import javax.crypto.spec.SecretKeySpec;
import javax.tools.JavaCompiler;
import javax.tools.JavaFileObject;
import java.awt.*;
import java.io.File;
import java.io.UnsupportedEncodingException;
import java.lang.ref.*;
import java.lang.reflect.Type;
import java.math.BigDecimal;
import java.net.Socket;
import java.net.URLDecoder;
import java.nio.ByteBuffer;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
@@ -1538,4 +1542,37 @@ public class T {
super(value);
}
}
/**
* 获取项目中各种路径
*
* @author ThinkPad
*/
public static class WebPathUtil {
private static final Log log = Log.get();
/**
* 如果已打成jar包则返回jar包所在目录
* 如果未打成jar则返回target所在目录
*
* @return
*/
public static String getClassPath() {
try {
// 项目的编译文件的根目录
String path = URLDecoder.decode(System.getProperty("user.dir"), "utf-8");
log.debug("root path:{}", path);
return path;
} catch (UnsupportedEncodingException e) {
return null;
}
}
public static String getRootPath() {
File file = T.FileUtil.file(WebPathUtil.getClassPath());
return file.getAbsolutePath();
}
}
}

View File

@@ -0,0 +1,169 @@
package net.geedge.asw.module.runner.controller;
import cn.dev33.satoken.annotation.SaIgnore;
import cn.hutool.core.lang.Opt;
import cn.hutool.log.Log;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import jakarta.servlet.http.HttpServletResponse;
import net.geedge.asw.common.util.R;
import net.geedge.asw.common.util.RCode;
import net.geedge.asw.common.util.T;
import net.geedge.asw.module.app.entity.PackageEntity;
import net.geedge.asw.module.runner.entity.JobEntity;
import net.geedge.asw.module.runner.entity.PlaybookEntity;
import net.geedge.asw.module.runner.entity.RunnerEntity;
import net.geedge.asw.module.runner.service.IJobService;
import net.geedge.asw.module.runner.service.IRunnerService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import java.io.IOException;
import java.util.Map;
@RestController
@RequestMapping("/api/v1/runner")
public class RunnerController {
private static final Log log = Log.get();
@Autowired
private IJobService jobService;
@Autowired
private IRunnerService runnerService;
@GetMapping("/{id}")
public R detail(@PathVariable("id") String id) {
RunnerEntity runnerEntity = runnerService.getById(id);
return R.ok().putData("record", runnerEntity);
}
@GetMapping
public R list(@RequestParam Map<String, Object> params) {
T.VerifyUtil.is(params).notNull()
.and(T.MapUtil.getStr(params, "workspaceId")).notEmpty(RCode.WORKSPACE_ID_CANNOT_EMPTY);
Page page = runnerService.queryList(params);
return R.ok(page);
}
@PostMapping
public R add(@RequestBody RunnerEntity entity) {
T.VerifyUtil.is(entity).notNull()
.and(entity.getWorkspaceId()).notEmpty(RCode.WORKSPACE_ID_CANNOT_EMPTY);
RunnerEntity runner = runnerService.saveRunner(entity);
return R.ok().putData("record", runner);
}
@PutMapping
public R update(@RequestBody RunnerEntity entity) {
T.VerifyUtil.is(entity).notNull()
.and(entity.getId()).notEmpty(RCode.ID_CANNOT_EMPTY)
.and(entity.getWorkspaceId()).notEmpty(RCode.WORKSPACE_ID_CANNOT_EMPTY);
RunnerEntity runner = runnerService.updateRunner(entity);
return R.ok().putData("record", runner);
}
@DeleteMapping("/{id}")
public R delete(@PathVariable("id") String id) {
runnerService.removeById(id);
return R.ok();
}
@SaIgnore
@PostMapping("/register")
public void register(@RequestHeader("Authorization") String token, HttpServletResponse response) throws IOException {
RunnerEntity runner = runnerService.getOne(new LambdaUpdateWrapper<RunnerEntity>().eq(RunnerEntity::getToken, token));
String status = Opt.ofNullable(runner).map(RunnerEntity::getStatus).orElseGet(() -> null);
if (!T.StrUtil.equals("online", status)) {
log.warn("[register] [runner is offline] [token: {}]", token);
response.sendError(HttpServletResponse.SC_FORBIDDEN, "Runner is offline");
}
}
@SaIgnore
@PostMapping("/heartbeat")
public void heartbeat(@RequestHeader("Authorization") String token, @RequestBody Map<String, Integer> platformMap,
HttpServletResponse response) throws IOException {
RunnerEntity runner = runnerService.getOne(new LambdaUpdateWrapper<RunnerEntity>().eq(RunnerEntity::getToken, token));
String status = Opt.ofNullable(runner).map(RunnerEntity::getStatus).orElseGet(() -> null);
if (!T.StrUtil.equals("online", status)) {
log.warn("[heartbeat] [runner is offline] [token: {}]", token);
response.sendError(HttpServletResponse.SC_FORBIDDEN, "Runner is offline");
return;
}
// findjob by platform
String platform = platformMap.entrySet().stream().filter(entry -> entry.getValue() > 0).findFirst().map(entry -> entry.getKey()).orElseGet(null);
JobEntity job = jobService.assignPendingJob(runner.getId(), platform);
if (T.ObjectUtil.isNotNull(job)) {
// package
PackageEntity pkg = job.getPkg();
Map<String, String> pkgInfo = T.MapUtil.builder("id", pkg.getId())
.put("platform", pkg.getPlatform())
.put("identifier", pkg.getIdentifier())
.put("version", pkg.getVersion())
.build();
// playbook
PlaybookEntity playbook = job.getPlaybook();
Map<String, String> pbInfo = T.MapUtil.builder("id", playbook.getId())
.put("name", playbook.getName())
.build();
// response job info
Map<Object, Object> responseData = T.MapUtil.builder()
.put("id", job.getId())
.put("pkg", pkgInfo)
.put("playbook", pbInfo)
.build();
response.setCharacterEncoding("UTF-8");
response.setContentType("text/html; charset=UTF-8");
response.getWriter().write(T.JSONUtil.toJsonStr(responseData));
}
}
@SaIgnore
@PutMapping("/trace/{jobId}")
public void trace(@RequestHeader("Authorization") String token, @PathVariable String jobId, @RequestBody byte[] bytes,
HttpServletResponse response) throws IOException {
RunnerEntity runner = runnerService.getOne(new LambdaUpdateWrapper<RunnerEntity>().eq(RunnerEntity::getToken, token));
String status = Opt.ofNullable(runner).map(RunnerEntity::getStatus).orElseGet(() -> null);
if (!T.StrUtil.equals("online", status)) {
log.warn("[trace] [runner is offline] [token: {}]", token);
response.sendError(HttpServletResponse.SC_FORBIDDEN, "Runner is offline");
return;
}
try {
// 追加到文件中
String content = T.StrUtil.str(bytes, T.CharsetUtil.CHARSET_UTF_8);
jobService.appendTraceLogStrToFile(jobId, content);
} catch (Exception e) {
log.error("[trace] [error] [job: {}]", jobId);
response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
}
}
@SaIgnore
@PutMapping("/jobResult/{jobId}")
public void jobResult(@RequestHeader("Authorization") String token, @PathVariable String jobId, @RequestParam String state,
@RequestParam(value = "file", required = false) MultipartFile pcapFile,
HttpServletResponse response) throws IOException {
RunnerEntity runner = runnerService.getOne(new LambdaUpdateWrapper<RunnerEntity>().eq(RunnerEntity::getToken, token));
String status = Opt.ofNullable(runner).map(RunnerEntity::getStatus).orElseGet(() -> null);
if (!T.StrUtil.equals("online", status)) {
log.warn("[trace] [runner is offline] [token: {}]", token);
response.sendError(HttpServletResponse.SC_FORBIDDEN, "Runner is offline");
return;
}
// 更新任务状态
jobService.updateJobResult(jobId, state, pcapFile);
}
}

View File

@@ -4,6 +4,7 @@ import com.baomidou.mybatisplus.core.mapper.BaseMapper;
import com.baomidou.mybatisplus.core.metadata.IPage;
import net.geedge.asw.module.runner.entity.JobEntity;
import org.apache.ibatis.annotations.Mapper;
import org.apache.ibatis.annotations.Param;
import java.util.List;
import java.util.Map;
@@ -13,4 +14,6 @@ public interface JobDao extends BaseMapper<JobEntity>{
List<JobEntity> queryList(IPage page, Map<String, Object> params);
JobEntity getPendingJobByPlatform(@Param("platform") String platform);
}

View File

@@ -1,10 +1,16 @@
package net.geedge.asw.module.runner.dao;
import com.baomidou.mybatisplus.core.mapper.BaseMapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import net.geedge.asw.module.runner.entity.RunnerEntity;
import org.apache.ibatis.annotations.Mapper;
import java.util.List;
import java.util.Map;
@Mapper
public interface RunnerDao extends BaseMapper<RunnerEntity>{
List<RunnerEntity> queryList(Page page, Map<String, Object> params);
}

View File

@@ -12,6 +12,7 @@ public class RunnerEntity {
@TableId(type = IdType.ASSIGN_UUID)
private String id;
private String name;
private String token;
private String tags;
private String supportPlatforms;
private Integer shareFlag;

View File

@@ -3,7 +3,9 @@ package net.geedge.asw.module.runner.service;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.IService;
import net.geedge.asw.module.runner.entity.JobEntity;
import org.springframework.web.multipart.MultipartFile;
import java.io.IOException;
import java.util.List;
import java.util.Map;
@@ -17,4 +19,10 @@ public interface IJobService extends IService<JobEntity>{
void removeJob(List<String> ids);
JobEntity assignPendingJob(String id, String platform);
void appendTraceLogStrToFile(String jobId, String content) throws RuntimeException;
void updateJobResult(String jobId, String state, MultipartFile pcapFile) throws IOException;
}

View File

@@ -3,6 +3,10 @@ package net.geedge.asw.module.runner.service;
import com.baomidou.mybatisplus.extension.service.IService;
import net.geedge.asw.module.runner.entity.PcapEntity;
import java.io.InputStream;
public interface IPcapService extends IService<PcapEntity>{
void savePcap(String jobId, InputStream inputStream);
}

View File

@@ -1,8 +1,17 @@
package net.geedge.asw.module.runner.service;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.IService;
import net.geedge.asw.module.runner.entity.RunnerEntity;
import java.util.Map;
public interface IRunnerService extends IService<RunnerEntity>{
Page queryList(Map<String, Object> params);
RunnerEntity saveRunner(RunnerEntity entity);
RunnerEntity updateRunner(RunnerEntity entity);
}

View File

@@ -1,6 +1,9 @@
package net.geedge.asw.module.runner.service.impl;
import cn.dev33.satoken.stp.StpUtil;
import cn.hutool.core.io.IORuntimeException;
import cn.hutool.log.Log;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import net.geedge.asw.common.util.RCode;
@@ -14,20 +17,30 @@ import net.geedge.asw.module.runner.entity.JobEntity;
import net.geedge.asw.module.runner.entity.PlaybookEntity;
import net.geedge.asw.module.runner.entity.RunnerEntity;
import net.geedge.asw.module.runner.service.IJobService;
import net.geedge.asw.module.runner.service.IPcapService;
import net.geedge.asw.module.runner.service.IPlaybookService;
import net.geedge.asw.module.runner.service.IRunnerService;
import net.geedge.asw.module.runner.util.RunnerConstant;
import net.geedge.asw.module.workbook.service.IWorkbookResourceService;
import net.geedge.asw.module.workbook.util.WorkbookConstant;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.multipart.MultipartFile;
import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.Map;
@Service
public class JobServiceImpl extends ServiceImpl<JobDao, JobEntity> implements IJobService {
private static final Log log = Log.get();
@Autowired
private IPcapService pcapService;
@Autowired
private IRunnerService runnerService;
@@ -43,6 +56,16 @@ public class JobServiceImpl extends ServiceImpl<JobDao, JobEntity> implements IJ
@Autowired
private IWorkbookResourceService workbookResourceService;
/**
* rootPath/result/{jobId}
*
* @param jobId
* @return
*/
private String getJobResultPath(String jobId) {
return T.FileUtil.file(T.WebPathUtil.getRootPath(), "result", jobId).getPath();
}
@Override
public JobEntity queryInfo(String id) {
JobEntity job = this.getById(id);
@@ -85,6 +108,12 @@ public class JobServiceImpl extends ServiceImpl<JobDao, JobEntity> implements IJ
// workbook resource
workbookResourceService.saveResource(entity.getWorkbookId(), entity.getId(), WorkbookConstant.ResourceType.JOB.getValue());
// trace log file path
File traceLogFile = T.FileUtil.file(this.getJobResultPath(entity.getId()), "trace.log");
this.update(new LambdaUpdateWrapper<JobEntity>()
.set(JobEntity::getLogPath, traceLogFile.getPath())
.eq(JobEntity::getId, entity.getId()));
return entity;
}
@@ -97,4 +126,56 @@ public class JobServiceImpl extends ServiceImpl<JobDao, JobEntity> implements IJ
workbookResourceService.removeResource(ids, WorkbookConstant.ResourceType.JOB.getValue());
}
@Override
public synchronized JobEntity assignPendingJob(String runnerId, String platform) {
if (T.StrUtil.hasEmpty(runnerId, platform)) {
return null;
}
// query
JobEntity job = this.getBaseMapper().getPendingJobByPlatform(platform);
if (T.ObjectUtil.isNotNull(job)) {
// update
this.update(new LambdaUpdateWrapper<JobEntity>()
.set(JobEntity::getRunnerId, runnerId)
.set(JobEntity::getStatus, RunnerConstant.JobStatus.RUNNING.getValue())
.set(JobEntity::getStartTimestamp, System.currentTimeMillis())
.eq(JobEntity::getId, job.getId())
);
}
return job;
}
@Override
public void appendTraceLogStrToFile(String jobId, String content) throws RuntimeException {
try {
JobEntity job = this.getById(jobId);
if (T.StrUtil.isEmpty(job.getLogPath())) {
File traceLogFile = T.FileUtil.file(this.getJobResultPath(jobId), "trace.log");
job.setLogPath(traceLogFile.getPath());
}
// append content
T.FileUtil.appendString(content, T.FileUtil.file(job.getLogPath()), T.CharsetUtil.CHARSET_UTF_8);
} catch (IORuntimeException e) {
log.error(e, "[appendTraceLogStrToFile] [error] [job: {}] [content: {}]", jobId, content);
throw new RuntimeException(e.getMessage());
}
}
@Override
@Transactional(rollbackFor = Exception.class)
public void updateJobResult(String jobId, String state, MultipartFile pcapFile) throws IOException {
// update job status
state = T.StrUtil.equals("success", state) ? RunnerConstant.JobStatus.PASSED.getValue() : state;
this.update(new LambdaUpdateWrapper<JobEntity>()
.set(JobEntity::getStatus, state)
.set(JobEntity::getEndTimestamp, System.currentTimeMillis())
.eq(JobEntity::getId, jobId)
);
// save pcap file
if (T.ObjectUtil.isNotNull(pcapFile)) {
pcapService.savePcap(jobId, pcapFile.getInputStream());
}
}
}

View File

@@ -6,8 +6,16 @@ import net.geedge.asw.module.runner.entity.PcapEntity;
import net.geedge.asw.module.runner.service.IPcapService;
import org.springframework.stereotype.Service;
import java.io.InputStream;
@Service
public class PcapServiceImpl extends ServiceImpl<PcapDao, PcapEntity> implements IPcapService {
@Override
public void savePcap(String jobId, InputStream inputStream) {
// TODO
}
}

View File

@@ -1,13 +1,51 @@
package net.geedge.asw.module.runner.service.impl;
import cn.dev33.satoken.stp.StpUtil;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import net.geedge.asw.common.util.T;
import net.geedge.asw.module.runner.dao.RunnerDao;
import net.geedge.asw.module.runner.entity.RunnerEntity;
import net.geedge.asw.module.runner.service.IRunnerService;
import org.springframework.stereotype.Service;
import java.util.List;
import java.util.Map;
@Service
public class RunnerServiceImpl extends ServiceImpl<RunnerDao, RunnerEntity> implements IRunnerService {
@Override
public Page queryList(Map<String, Object> params) {
Page page = T.PageUtil.getPage(params);
List<RunnerEntity> jobList = this.getBaseMapper().queryList(page, params);
page.setRecords(jobList);
return page;
}
}
@Override
public RunnerEntity saveRunner(RunnerEntity entity) {
entity.setCreateTimestamp(System.currentTimeMillis());
entity.setUpdateTimestamp(System.currentTimeMillis());
entity.setCreateUserId(StpUtil.getLoginIdAsString());
entity.setUpdateUserId(StpUtil.getLoginIdAsString());
// token
entity.setToken(T.IdUtil.fastSimpleUUID());
// save
this.save(entity);
return entity;
}
@Override
public RunnerEntity updateRunner(RunnerEntity entity) {
entity.setUpdateTimestamp(System.currentTimeMillis());
entity.setUpdateUserId(StpUtil.getLoginIdAsString());
// update
this.updateById(entity);
return entity;
}
}

View File

@@ -0,0 +1,54 @@
package net.geedge.asw.module.runner.util;
public class RunnerConstant {
/**
* job status
*/
public enum JobStatus {
CREATED("created"),
PENDING("pending"),
RUNNING("running"),
PASSED("passed"),
FAILED("failed"),
CANCEL("cancel");
private String value;
JobStatus(String value) {
this.value = value;
}
public String getValue() {
return value;
}
}
/**
* pcap status
*/
public enum PcapStatus {
UPLOADED("Uploaded"),
ANALYZING("Analyzing"),
COMPLETED("Completed");
private String value;
PcapStatus(String value) {
this.value = value;
}
public String getValue() {
return value;
}
}
}

View File

@@ -26,6 +26,7 @@
<association property="pkg" columnPrefix="pkg_" javaType="net.geedge.asw.module.app.entity.PackageEntity">
<id property="id" column="id"/>
<result property="platform" column="platform"/>
<result property="identifier" column="identifier"/>
<result property="version" column="version"/>
<result property="logo" column="logo"/>
</association>
@@ -56,6 +57,7 @@
pkg.platform AS pkg_platform,
pkg.version AS pkg_version,
pkg.logo AS pkg_logo,
pkg.identifier AS pkg_identifier,
app.id AS app_id,
app.name AS app_name,
@@ -121,4 +123,26 @@
</if>
</select>
<select id="getPendingJobByPlatform" resultMap="jobResultMap">
SELECT
job.*,
pkg.id AS pkg_id,
pkg.platform AS pkg_platform,
pkg.identifier AS pkg_identifier,
pkg.version AS pkg_version,
pb.id AS pb_id,
pb.name AS pb_name
FROM
job job
LEFT JOIN package pkg ON job.package_id = pkg.id
LEFT JOIN playbook pb ON job.playbook_id = pb.id
WHERE
job.status = 'pending' and pkg.platform = #{platform}
ORDER BY job.create_timestamp ASC
LIMIT 1
</select>
</mapper>

View File

@@ -0,0 +1,32 @@
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE mapper PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN" "http://mybatis.org/dtd/mybatis-3-mapper.dtd">
<mapper namespace="net.geedge.asw.module.runner.dao.RunnerDao">
<select id="queryList" resultType="net.geedge.asw.module.runner.entity.RunnerEntity">
SELECT
*
FROM
runner
<where>
<if test="params.workspaceId != null and params.workspaceId != ''">
workspace_id = #{params.workspaceId}
</if>
<if test="params.q != null and params.q != ''">
AND locate(#{params.q}, description)
</if>
<if test="params.tags != null and params.tags != ''">
AND <foreach item="item" collection="params.tags.split(',')" separator="OR" index="" open="(" close=")">
locate(#{item}, tags)
</foreach>
</if>
</where>
<if test="params.orderBy == null or params.orderBy == ''">
ORDER BY id
</if>
</select>
</mapper>

View File

@@ -141,6 +141,7 @@ DROP TABLE IF EXISTS `runner`;
CREATE TABLE `runner` (
`id` varchar(64) NOT NULL COMMENT '主键',
`name` varchar(256) NOT NULL DEFAULT '' COMMENT '名称',
`token` varchar(64) NOT NULL DEFAULT '' COMMENT 'token',
`tags` varchar(256) NOT NULL DEFAULT '' COMMENT '标签,多个逗号分隔',
`support_platforms` varchar(256) NOT NULL DEFAULT '' COMMENT '支持的平台; 可选值android,ios,windows; 多个逗号分隔; 例android,ios',
`share_flag` int(1) NOT NULL DEFAULT 1 COMMENT '共享标识; 1共享 0不共享仅创建人可用',
@@ -192,8 +193,8 @@ CREATE TABLE `job` (
`schedule_id` varchar(64) NOT NULL DEFAULT '' COMMENT '定时器ID',
`signature_ids` text NOT NULL DEFAULT '' COMMENT '特征ID,多个逗号分隔',
`tags` varchar(256) NOT NULL DEFAULT '' COMMENT '标签; 默认:"";多个用逗号分隔;例kz,vpn,android',
`start_timestamp` bigint(20) NOT NULL DEFAULT (UNIX_TIMESTAMP(NOW()) * 1000) COMMENT '开始时间戳',
`end_timestamp` bigint(20) NOT NULL DEFAULT (UNIX_TIMESTAMP(NOW()) * 1000) COMMENT '结束时间戳',
`start_timestamp` bigint(20) NOT NULL DEFAULT -1 COMMENT '开始时间戳',
`end_timestamp` bigint(20) NOT NULL DEFAULT -1 COMMENT '结束时间戳',
`status` varchar(64) NOT NULL DEFAULT '' COMMENT '状态; 可选值: created,pending,running,passed,failed,cancel',
`pcap_id` varchar(64) NOT NULL DEFAULT '' COMMENT 'PCAP ID',
`log_path` varchar(256) NOT NULL DEFAULT '' COMMENT '日志文件路径',