feat: NEZ-2021 ota接口开发

4 files changed, 157 insertions, 1 deletions

diff --git a/src/main/java/net/geedge/confagent/controller/BaseController.java b/src/main/java/net/geedge/confagent/controller/BaseController.java
index a92f204..1ce7266 100644
--- a/src/main/java/net/geedge/confagent/controller/BaseController.java
+++ b/src/main/java/net/geedge/confagent/controller/BaseController.java
@@ -1,5 +1,6 @@
 package net.geedge.confagent.controller;
 
+import net.geedge.confagent.util.Tool;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.http.HttpEntity;
@@ -47,6 +48,10 @@ public abstract class BaseController {
 	protected static final int DEFAULT_SNMPEXPORTER_PORT = 9116; // snmp exporter 默认监听端口
 
 	protected static final int DEFAULT_BLACKBOXEXPORTER_PORT = 9115; // black box exporter 默认监听端口
+	/**
+	 * aes 对称加密 密钥
+	 */
+	public static final byte[] AES_SECRET_KEY = Tool.HexUtil.decodeHex("bde5430614b21baf1c53bd6f616d1a39");
 
 	@Value("${confagent.prometheus.config}")
 	protected String prometheusConfPath;
diff --git a/src/main/java/net/geedge/confagent/controller/OTAController.java b/src/main/java/net/geedge/confagent/controller/OTAController.java
new file mode 100644
index 0000000..71636c5
--- /dev/null
+++ b/src/main/java/net/geedge/confagent/controller/OTAController.java
@@ -0,0 +1,79 @@
+package net.geedge.confagent.controller;
+
+import cn.hutool.crypto.digest.MD5;
+import cn.hutool.log.Log;
+import net.geedge.confagent.util.*;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.web.bind.annotation.PostMapping;
+import org.springframework.web.bind.annotation.RequestParam;
+import org.springframework.web.bind.annotation.RestController;
+import org.springframework.web.multipart.MultipartFile;
+
+import javax.servlet.http.HttpServletRequest;
+import java.io.File;
+
+@RestController
+public class OTAController extends BaseController{
+
+    private static Log log = Log.get();
+
+    @Autowired
+    private ConfagentUtil confagentUtil;
+
+    @Value("${version}")
+    private String currentVersion;
+
+    private String rpmPackagePath = "/tmp/";
+
+    @PostMapping("/ota")
+    public R ota(HttpServletRequest request, @RequestParam(name = "file", required = false) MultipartFile file, String md5){
+
+        //检查token
+        String token = request.getHeader("Authorization");
+        if (Tool.StrUtil.isBlank(token)) {
+            return R.error(RCode.AUTH_TOKEN_ISNULL);
+        }else {
+            if (confagentUtil.checkToken(token).getCode() != RCode.SUCCESS.getCode())return R.error(RCode.AUTH_TOKEN_INVALID);
+        }
+
+        //缓存文件
+        String filePath = rpmPackagePath + String.format("nezha.%s.rpm", CommonUtils.uuid());
+        log.info("cache rpm package name : {}", filePath);
+        File dest = new File(filePath);
+        try {
+            file.transferTo(dest);
+        } catch (Exception e) {
+            return R.error(RCode.OTA_FILE_CACHE_ERROR);
+        }
+
+        //校验MD5
+        String updateFileMd5 = MD5.create().digestHex(dest);
+        md5 = Tool.AesUtil.decrypt(md5, AES_SECRET_KEY);
+        if(!updateFileMd5.equals(md5)){
+            return R.error(RCode.OTA_MD5_ERROR);
+        }
+
+        //校验版本号
+        String fileName = file.getOriginalFilename();
+        String updateVersion = fileName.split("-")[2];
+        updateVersion = updateVersion.substring(0,updateVersion.lastIndexOf("."));
+        currentVersion = currentVersion.substring(0,currentVersion.lastIndexOf("."));
+        log.info("current version : {} , update version : {}", currentVersion , updateVersion);
+        if( !(Tool.DateUtil.betweenDay(Tool.DateUtil.parseDate(updateVersion),Tool.DateUtil.parseDate(currentVersion), false) > 0)){
+            return R.error(RCode.OTA_VERSION_ERROR);
+        }
+        //更新
+        String updateCommand = String.format("rpm -Uvh %s", filePath);
+        log.info("update nz-agent : {}", updateCommand);
+        new Thread(() -> {
+            try {
+                Tool.RuntimeUtil.exec(updateCommand);
+            } catch (Exception e) {
+                log.error("update rpm package error : {}", e.getMessage());
+            }
+        }).start();
+        return R.ok();
+    }
+
+}
diff --git a/src/main/java/net/geedge/confagent/util/RCode.java b/src/main/java/net/geedge/confagent/util/RCode.java
index ff16cc3..8e40b58 100644
--- a/src/main/java/net/geedge/confagent/util/RCode.java
+++ b/src/main/java/net/geedge/confagent/util/RCode.java
@@ -23,7 +23,11 @@ public enum RCode {
 	STOP_CMD_ERROR(10112,"The stop command is error"),
 
 	IP_ADDRESS_ISNULL(10113, "IP address can not be empty"),
-	
+
+	OTA_FILE_CACHE_ERROR(10114, "Cache rpm package failed"),
+	OTA_MD5_ERROR(10115, "MD5 is error"),
+	OTA_VERSION_ERROR(10116, "The updated version is lower than the current version"),
+
     ERROR(999, "error"); //通用错误/未知错误
 	private RCode(Integer code, String msg) {
 		this.code = code;
diff --git a/src/main/java/net/geedge/confagent/util/Tool.java b/src/main/java/net/geedge/confagent/util/Tool.java
index a89f174..dc092c1 100644
--- a/src/main/java/net/geedge/confagent/util/Tool.java
+++ b/src/main/java/net/geedge/confagent/util/Tool.java
@@ -16,11 +16,17 @@ import java.lang.reflect.Type;
 import java.math.BigDecimal;
 import java.net.*;
 import java.nio.ByteBuffer;
+import java.security.NoSuchAlgorithmException;
+import java.security.SecureRandom;
 import java.time.LocalDateTime;
 import java.time.temporal.Temporal;
 import java.time.temporal.TemporalAccessor;
 import java.util.*;
 
+import javax.crypto.Cipher;
+import javax.crypto.KeyGenerator;
+import javax.crypto.SecretKey;
+import javax.crypto.spec.SecretKeySpec;
 import javax.tools.JavaCompiler;
 import javax.tools.JavaFileObject;
 
@@ -1146,4 +1152,66 @@ public class Tool {
 			}
 		}
 	}
+	//aes 工具
+	public static class AesUtil {
+
+		private static final String KEY_ALGORITHM = "AES";
+		private static final String DEFAULT_CIPHER_ALGORITHM = "AES/ECB/PKCS5Padding";// 默认的加密算法
+
+		public static String encrypt(String content, byte[] key) {
+			try {
+				Cipher cipher = Cipher.getInstance(DEFAULT_CIPHER_ALGORITHM);// 创建密码器
+
+				byte[] byteContent = content.getBytes("utf-8");
+
+				cipher.init(Cipher.ENCRYPT_MODE, getSecretKey(key));// 初始化为加密模式的密码器
+
+				byte[] result = cipher.doFinal(byteContent);// 加密
+
+				return cn.hutool.core.util.HexUtil.encodeHexStr(result);
+			} catch (Exception ex) {
+				ex.printStackTrace();
+			}
+
+			return null;
+		}
+
+		public static String decrypt(String content, byte[] key) {
+
+			try {
+				// 实例化
+				Cipher cipher = Cipher.getInstance(DEFAULT_CIPHER_ALGORITHM);
+
+				// 使用密钥初始化，设置为解密模式
+				cipher.init(Cipher.DECRYPT_MODE, getSecretKey(key));
+
+				// 执行操作
+				byte[] result = cipher.doFinal(cn.hutool.core.util.HexUtil.decodeHex(content));
+				String s = new String(result, "utf-8");
+				return s;
+			} catch (Exception ex) {
+				ex.printStackTrace();
+			}
+			return null;
+		}
+
+		private static SecretKeySpec getSecretKey(byte[] key) {
+			// 返回生成指定算法密钥生成器的 KeyGenerator 对象
+			KeyGenerator kg = null;
+			try {
+				kg = KeyGenerator.getInstance(KEY_ALGORITHM);
+				SecureRandom random = SecureRandom.getInstance("SHA1PRNG");
+				random.setSeed(key);
+				// AES 要求密钥长度为 128
+				kg.init(128, random);
+				// 生成一个密钥
+				SecretKey secretKey = kg.generateKey();
+				return new SecretKeySpec(secretKey.getEncoded(), KEY_ALGORITHM);// 转换为AES专用密钥
+			} catch (NoSuchAlgorithmException ex) {
+				ex.printStackTrace();
+			}
+			return null;
+		}
+
+	}
 }