fix: 优化Agentcontroller代码

3 files changed, 169 insertions, 169 deletions

diff --git a/nz-admin/src/main/java/com/nis/common/utils/ResponseUtil.java b/nz-admin/src/main/java/com/nis/common/utils/ResponseUtil.java
new file mode 100644
index 00000000..dab14a30
--- /dev/null
+++ b/nz-admin/src/main/java/com/nis/common/utils/ResponseUtil.java
@@ -0,0 +1,22 @@
+package com.nis.common.utils;
+
+import java.io.File;
+import java.io.IOException;
+import javax.servlet.http.HttpServletResponse;
+
+import org.springframework.http.MediaType;
+
+import cn.hutool.core.io.IORuntimeException;
+
+public class ResponseUtil {
+
+	public static void downloadFile(HttpServletResponse response,File file) throws IORuntimeException, IOException {
+		response.reset();
+		response.setContentType(MediaType.APPLICATION_OCTET_STREAM.toString());
+		String fileName = Tool.URLUtil.encode(file.getName(), Tool.CharsetUtil.UTF_8);
+		response.addHeader("Content-Disposition", "attachment; filename=" + fileName);
+		response.addHeader("Content-Length", "" + file.length());
+		Tool.FileUtil.writeToStream(file, response.getOutputStream());
+	}
+	
+}
diff --git a/nz-admin/src/main/java/com/nis/common/utils/Tool.java b/nz-admin/src/main/java/com/nis/common/utils/Tool.java
index 8ff696ff..1ce3cbe4 100644
--- a/nz-admin/src/main/java/com/nis/common/utils/Tool.java
+++ b/nz-admin/src/main/java/com/nis/common/utils/Tool.java
@@ -1057,5 +1057,10 @@ public class Tool {
 	 */
 	public static class JSONUtil extends com.nis.common.utils.JSONUtil {};
 	
-
+	/**
+	 * http response 工具类
+	 * @author ThinkPad
+	 *
+	 */
+	public static class ResponseUtil extends com.nis.common.utils.ResponseUtil{};
 }
diff --git a/nz-admin/src/main/java/com/nis/modules/agent/AgentController.java b/nz-admin/src/main/java/com/nis/modules/agent/AgentController.java
index eeb47b91..51478344 100644
--- a/nz-admin/src/main/java/com/nis/modules/agent/AgentController.java
+++ b/nz-admin/src/main/java/com/nis/modules/agent/AgentController.java
@@ -1,182 +1,155 @@
 package com.nis.modules.agent;
 
-import cn.hutool.core.net.url.UrlBuilder;
-import cn.hutool.core.util.ReflectUtil;
-import cn.hutool.http.HttpConnection;
-import cn.hutool.http.HttpRequest;
-import cn.hutool.http.HttpResponse;
-import cn.hutool.json.JSONObject;
-import cn.hutool.log.Log;
+import java.io.File;
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.ui.freemarker.FreeMarkerTemplateUtils;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.PathVariable;
+import org.springframework.web.bind.annotation.PostMapping;
+import org.springframework.web.bind.annotation.RequestBody;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RequestParam;
+import org.springframework.web.bind.annotation.RestController;
+
+import com.nis.common.exception.NZException;
 import com.nis.common.smartvalidate.ValidateUtils;
-import com.nis.common.utils.*;
+import com.nis.common.utils.Constant;
+import com.nis.common.utils.R;
+import com.nis.common.utils.RCode;
+import com.nis.common.utils.TemplateUtil;
+import com.nis.common.utils.Tool;
 import com.nis.modules.promserver.entity.Promserver;
 import com.nis.modules.promserver.service.PromserverService;
 import com.nis.modules.sys.service.SysConfigService;
+
+import cn.hutool.core.io.IORuntimeException;
+import cn.hutool.core.net.url.UrlBuilder;
+import cn.hutool.json.JSONObject;
+import cn.hutool.log.Log;
 import freemarker.template.Template;
 import freemarker.template.TemplateException;
-import org.springframework.beans.factory.annotation.Autowired;
-import org.springframework.beans.factory.annotation.Value;
-import org.springframework.ui.freemarker.FreeMarkerTemplateUtils;
-import org.springframework.web.bind.annotation.*;
-
-import javax.servlet.http.HttpServletRequest;
-import javax.servlet.http.HttpServletResponse;
-import java.io.*;
-import java.net.HttpURLConnection;
-import java.util.HashMap;
-import java.util.Map;
 
 @RestController
 @RequestMapping("/agent")
 public class AgentController {
-    private Log log = Log.get();
-
-    @Value("${nezha.resourcePath:./}")
-    private String resourcePath;
-
-    @Autowired
-    private PromserverService promserverService;
-
-    @Autowired
-    private SysConfigService sysConfigService;
-
-    /**
-    * @Description 下载 install.sh脚本
-    * @Author  rui
-    * @Date   2021/4/2
-    */
-    @GetMapping("/{dcId}/{type}/install.sh")
-    public String pushInstallScript(HttpServletRequest request, @PathVariable("dcId") Integer dcId, @PathVariable("type") Integer type){
-        ValidateUtils.is(dcId).notNull(RCode.AGENT_INSTALL_PARAM_MISS).and(type).notNull(RCode.AGENT_INSTALL_PARAM_MISS);
-
-        String scriptPath = resourcePath+File.separator+"install"+ File.separator+"install.sh.template";
-
-        String scriptContent = Tool.FileUtil.readUtf8String(scriptPath);
-
-        String localhost = request.getLocalAddr(); //web component ip
-        Integer localPort = request.getLocalPort();
-
-        String remoteHost = request.getRemoteHost(); // agent component ip
-
-        String webToken = request.getHeader(Constant.AUTH_TOKEN_CODE);
-
-        Map<String,String> map = new HashMap<>();
-
-        map.put("remoteIP",localhost); // to install.sh is remote
-        map.put("remotePort",localPort+"");
-        map.put("localhost",remoteHost); // to install.sh is local
-        map.put("webToken",webToken);
-        map.put("datacenterID",dcId+"");
-        map.put("type",type+"");
-
-        scriptContent = replaceParamter(map,scriptContent);
-
-        return scriptContent;
-    }
-
-
-    /**
-    * @Description 下载confagent 的安装包
-    * @Author  rui
-    * @Date   2021/4/2
-    */
-    @GetMapping("/download")
-    public void downloadAgent(HttpServletResponse response, @RequestParam(defaultValue = "centos",required = false) String os){
-        log.info("download agent package to {}",os);
-        String json = sysConfigService.getValue("confagent_name");
-        JSONObject confagentNames = Tool.JSONUtil.parseObj(json);
-        String confagentName = (String)confagentNames.get(os);
-        downloadAgentToOs(response,confagentName);
-    }
-
-    /**
-    * @Description 将confagent自动注册 为promserver
-    * @Author  rui
-    * @Date   2021/4/2
-    */
-    @PostMapping("/register")
-    public R autoRegister(HttpServletRequest request,@RequestBody Promserver promserver){
-        if(Tool.StrUtil.isEmpty(promserver.getHost())){
-            promserver.setHost(request.getRemoteAddr());
-        }
-
-        return R.ok().putData("id",promserverService.saveOrUpdatePromserver(promserver, false));
-    }
-
-
-    /**
-    * @Description 访问confagent refresh接口，刷新confagent token
-    * @Author  rui
-    * @Date   2021/4/2
-    */
-    @PostMapping("/token/refresh")
-    public void refreshToken(HttpServletResponse response,@RequestBody Map<String,Object> params){
-        String host = (String)params.get("host");
-        String token = (String)params.get("token");
-        Integer port = (Integer)params.get("port");
-
-        ValidateUtils.is(host).notNull(RCode.AGENT_INSTALL_PARAM_MISS).and(token).notNull(RCode.AGENT_INSTALL_PARAM_MISS).and(port).notNull(RCode.AGENT_INSTALL_PARAM_MISS);
-
-        UrlBuilder urlBuilder = UrlBuilder.create();
-        urlBuilder.setScheme("http").setHost(host).setPort(port).appendPath("confagent").appendPath("auth").appendPath("refresh");
-        refreshToken(urlBuilder.build(),token,response);
-    }
-
-    private String replaceParamter(Map<String,String> keyValue,String content){
-        if(Tool.MapUtil.isEmpty(keyValue)) return content;
-
-        try{
-            Template template = TemplateUtil.stringToTemplate(content,"installTemplate");
-
-            content = FreeMarkerTemplateUtils.processTemplateIntoString(template, keyValue);
-
-        }catch (IOException | TemplateException e){
-            log.error("install template parse error",e);
-        }
-        return content;
-    }
-
-    private void downloadAgentToOs(HttpServletResponse response, String filename){
-        File file = Tool.FileUtil.file(resourcePath,"install",filename);
-        response.reset();
-        response.setContentType("application/octet-stream");
-        String fileName = Tool.URLUtil.encode(file.getName(), Tool.CharsetUtil.UTF_8);
-        response.addHeader("Content-Disposition", "attachment; filename=" +fileName );
-        response.addHeader("Content-Length", "" + file.length());
-
-        try {
-            Tool.FileUtil.writeToStream(file,new BufferedOutputStream(response.getOutputStream()));
-        } catch (IOException e) {
-            log.error("download agent package failed",e);
-        }
-    }
-
-    private void refreshToken(String url,String token,HttpServletResponse response){
-        OutputStream responseOutputStream = null;
-        InputStream connInputStream = null;
-        try{
-
-            response.setContentType("application/json");
-            responseOutputStream = response.getOutputStream();
-
-            HttpRequest get = Tool.HttpUtil.createGet(url);
-            get.auth(token);
-            HttpResponse executeResponse = get.execute(false);
-            connInputStream = executeResponse.bodyStream();
-            Tool.IoUtil.copy(connInputStream, responseOutputStream);
-            responseOutputStream.flush();//flush 输出流
-
-        }catch (Exception e){
-            try {
-                response.sendError(500, "request error");
-            } catch (IOException e1) {
-                log.error("proxy request error",e1);
-            }
-            log.error("request error : ",e);
-        }finally {
-            Tool.IoUtil.close(responseOutputStream);
-            Tool.IoUtil.close(connInputStream);
-        }
-    }
+	private Log log = Log.get();
+
+	@Value("${nezha.installFilePath:./install}")
+	private String installFilePath;
+
+	@Autowired
+	private PromserverService promserverService;
+
+	@Autowired
+	private SysConfigService sysConfigService;
+
+	/**
+	 * @Description 下载 install.sh脚本
+	 * @Author rui
+	 * @Date 2021/4/2
+	 */
+	@GetMapping("/{dcId}/{type}/install.sh")
+	public String pushInstallScript(HttpServletRequest request, @PathVariable("dcId") Integer dcId,
+			@PathVariable("type") Integer type) {
+		ValidateUtils.is(dcId).notNull(RCode.AGENT_INSTALL_PARAM_MISS).and(type)
+				.notNull(RCode.AGENT_INSTALL_PARAM_MISS);
+		// 读取sh脚本内容
+		File shFile = Tool.FileUtil.file(installFilePath, "install.sh.template");
+		// 组织变量内容
+		String scriptContent = Tool.FileUtil.readUtf8String(shFile);
+		String localhost = request.getLocalAddr(); // web component ip
+		Integer localPort = request.getLocalPort();
+		String remoteHost = request.getRemoteHost(); // agent component ip
+		String webToken = request.getHeader(Constant.AUTH_TOKEN_CODE);
+		Map<String, String> map = new HashMap<>();
+		map.put("remoteIP", localhost); // to install.sh is remote
+		map.put("remotePort", String.valueOf(localPort));
+		map.put("localhost", remoteHost); // to install.sh is local
+		map.put("webToken", webToken);
+		map.put("datacenterID",String.valueOf(dcId));
+		map.put("type", String.valueOf(type));
+		try {
+			// 根据模板文件生成 sh文件
+			Template template = TemplateUtil.stringToTemplate(scriptContent, "installTemplate");
+			scriptContent = FreeMarkerTemplateUtils.processTemplateIntoString(template, map);
+		} catch (IOException | TemplateException e) {
+			log.error("install template parse error", e);
+			return "echo 'Internal server error'";
+		}
+		return scriptContent;
+	}
+
+	/**
+	 * @Description 下载confagent 的安装包
+	 * @Author rui
+	 * @Date 2021/4/2
+	 */
+	@GetMapping("/download")
+	public void downloadAgent(HttpServletResponse response,
+			@RequestParam(defaultValue = "centos", required = false) String os) {
+		log.info("download agent package to {}", os);
+		String json = sysConfigService.getValue("confagent_name");
+		JSONObject confagentNames = Tool.JSONUtil.parseObj(json);
+		String confagentName = (String) confagentNames.get(os);
+		File file = Tool.FileUtil.file(installFilePath, confagentName);
+		log.info("file exist: {}", file.exists());
+		try {
+			Tool.ResponseUtil.downloadFile(response, file);
+		} catch (IORuntimeException | IOException e) {
+			log.error("file download error", e);
+			throw new NZException("file download error", 500);
+		}
+	}
+
+	/**
+	 * @Description 将confagent自动注册 为promserver
+	 * @Author rui
+	 * @Date 2021/4/2
+	 */
+	@PostMapping("/register")
+	public R autoRegister(HttpServletRequest request, @RequestBody Promserver promserver) {
+		if (Tool.StrUtil.isEmpty(promserver.getHost())) {
+			promserver.setHost(request.getRemoteAddr());
+		}
+		//保存 agent 
+		Integer id = promserverService.saveOrUpdatePromserver(promserver, false);
+		return R.ok().putData("id", id);
+	}
+
+	/**
+	 * @Description 访问confagent refresh接口，刷新confagent token
+	 * @Author rui
+	 * @Date 2021/4/2
+	 */
+	@PostMapping("/token/refresh")
+	public R refreshToken(HttpServletResponse response, @RequestBody Map<String, Object> params) {
+		String host = (String) params.get("host");
+		String token = (String) params.get("token");
+		Integer port = (Integer) params.get("port");
+		//校验参数
+		ValidateUtils.is(host).notNull(RCode.AGENT_INSTALL_PARAM_MISS).and(token)
+				.notNull(RCode.AGENT_INSTALL_PARAM_MISS).and(port).notNull(RCode.AGENT_INSTALL_PARAM_MISS);
+		//拼接url
+		UrlBuilder urlBuilder = UrlBuilder.create();
+		urlBuilder.setScheme("http").setHost(host).setPort(port).appendPath("confagent").appendPath("auth")
+				.appendPath("refresh");
+		//调用 confagent接口 获取结果
+		String result = Tool.HttpUtil.createGet(urlBuilder.build()).auth(token).execute().body();
+		log.info("result:{}", result);
+		//响应结果
+		if (Tool.StrUtil.isNotEmpty(result)) {
+			return Tool.JSONUtil.toBean(result, R.class);
+		} else {
+			return R.error(500, "token refresh failed");
+		}
+	}
 
 }