nezha-nz-talon/src/main/java/net/geedge/confagent/controller/PromtailController.java

package net.geedge.confagent.controller;

import cn.hutool.core.net.url.UrlBuilder;
import cn.hutool.core.net.url.UrlPath;
import cn.hutool.core.util.ReflectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HttpConnection;
import cn.hutool.log.Log;
import net.geedge.confagent.annotation.UnCheckToken;
import net.geedge.confagent.util.ConfagentUtil;
import net.geedge.confagent.util.R;
import net.geedge.confagent.util.RCode;
import net.geedge.confagent.util.Tool;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.bind.annotation.*;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.nio.charset.Charset;
import java.util.*;

import javax.servlet.ServletInputStream;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

@RestController
@RequestMapping("/promtail")
public class PromtailController extends BaseController{
    private final static Log log = Log.get();

    @Autowired
	private ConfagentUtil confagentUtil;
    
    @Value("${confagent.promtail.query.auth:true}")
    private Boolean queryAuth;

    @Value("${confagent.promtail.restart:systemctl restart promtail}")
    private String restartCmd;
    
    @Value("${confagent.versionFile:promtail.version}")
    private String versionFile;
    
    @Value("${confagent.promtail.startCmd:systemctl start promtail.service}")
    private String startCmd;
    
    @Value("${confagent.promtail.stopCmd:systemctl stop promtail.service}")
    private String stopCmd;
    
    private final String[] QUERY_API_SUFFIX = {"query","query_range","series","labels","label","values"};
    
    private static String rootPath = Tool.WebPathUtil.getRootPath();

    /**
    * @Description 获取promtail相关配置
    * @Author  rui
    * @Date   2021/3/24
    */
    @GetMapping("/config")
    public R queryConfig(){
        Map<String,String> cmdLine = ConfagentUtil.loadServiceConfigFile(promtailCmdLinePath);
        Map<String,Object> promtailConf = Tool.YamlUtil.readAsMap(promtailConfPath);
        String version = confagentUtil.readVersion();        
        
        Map<String,Object> result = new LinkedHashMap<>();
        
        result.put("version",version);     
        result.put("cmdline",cmdLine);
        result.put("config",promtailConf);

        return R.ok(result);
    }

    /**
    * @Description 写入promtail配置文件
    * @Author  rui
    * @Date   2021/3/25
    */
    @PostMapping("/config")
    public R overwriteConfig( @RequestBody Map<String,Object> configs){

    	Object version = configs.get("version");
    	File tf = Tool.FileUtil.file(rootPath, versionFile);
        log.debug("version file path : {}" ,tf.getAbsolutePath());
        if(version==null){
            return R.error(RCode.PROMTAIL_CONFIG_VERSION_ISNULL);
        }else {
        	Tool.FileUtil.writeUtf8String(version.toString(), tf);
        }
        
        Map<String,String> cmdLine =(Map<String,String>) configs.get("cmdline");
        Map<String,Object> promtailConf =(Map<String,Object>) configs.get("config");

        if(!Tool.MapUtil.isEmpty(cmdLine)){
            log.info("write promtail cmdLine conf:{}", Tool.JSONUtil.toJsonStr(cmdLine));
            writeServiceConfigFile(cmdLine,promtailCmdLinePath);
        }
        if(!Tool.MapUtil.isEmpty(promtailConf)){
            log.info("write promtail conf:{}", Tool.JSONUtil.toJsonStr(promtailConf));
            Tool.YamlUtil.writeAsMap(promtailConf,promtailConfPath);
        }
        
        //重启服务promtail
        if(Tool.StrUtil.isNotBlank(stopCmd)) {
        	log.info("stop promtail:"+stopCmd);
        	try {
        		Tool.RuntimeUtil.exec(stopCmd);
        	}catch(Exception e) {
        		log.error(e);
        		return R.error(RCode.PROMTAIL_STOP_CMD_ERROR);
        	}
        	log.info("stop promtail:"+stopCmd+"  end");
        }
        Tool.ThreadUtil.sleep(1000);
        if(Tool.StrUtil.isNotBlank(startCmd)) {
        	log.info("start promtail:"+startCmd);
        	try {
	//        	String[] b={"sh","-c",startCmd};
	//        	Tool.RuntimeUtil.exec(b);
	        	Tool.RuntimeUtil.exec(startCmd);
        	}catch(Exception e) {
        		log.error(e);
        		return R.error(RCode.PROMTAIL_START_CMD_ERROR);
        	}
        	log.info("start promtail:"+startCmd+"  end");
        }

        return R.ok();
    }
    
    /**
    * @Description 获取promtail相关配置版本
    * @Author  rui
    * @Date   2021/3/24
    */
    @GetMapping("/config/version")
    public R queryVersion(){
        Map<String,Object> result = new HashMap<>();        
        String version = confagentUtil.readVersion();        
        result.put("version",version);        
        return R.ok(result);
    }

    /**
    * @Description 写 promtail config.conf 启动参数配置文件
    * @Author  rui
    * @Date   2021/3/25
    */
    private void writeServiceConfigFile(Map<String,String> conf,String path){
        StringBuffer sb = new StringBuffer();
        sb.append("OPTION=\"");
        for (Map.Entry<String,String> entry:conf.entrySet()){
            String key = entry.getKey();
            String value = entry.getValue();
            sb.append("--"+key+"=");
            sb.append("'"+value+"' ");
        }
        sb.append("\"");

        Tool.FileUtil.writeUtf8String(sb.toString(),path);

    }

    /**
     * @Description 代理本promtail 接口
     * @Author  han
     * @Date   2021/8/10
     */
     @RequestMapping("/proxy/**")
     @UnCheckToken
     public void proxy(HttpServletRequest request, HttpServletResponse response){
         String promtailPath = request.getServletPath().replace("/promtail/proxy","");
         String token = request.getHeader("Authorization");

         R r = confagentUtil.checkToken(token);

         Boolean isQuery=false;

         //queryAuth 配置只限制查询是否需要校验
         if(queryAuth){
             isQuery = Tool.StrUtil.isNotBlank(Arrays.stream(QUERY_API_SUFFIX).filter(t -> promtailPath.indexOf(t) != -1).findAny().orElse(null));
         }

         if(isQuery &&r.getCode() != RCode.SUCCESS.getCode()){
             ConfagentUtil.writeResponse(response,r);
             return;
         }

         int port = 9080;
         Map<String,Object> promtailMap= Tool.YamlUtil.readAsMap(promtailConfPath);
         Map<String,Object> serverObj = (Map<String,Object>)promtailMap.get("server");
         if(serverObj!=null) {
         	port = (Integer)serverObj.get("http_listen_port");
         }

         requestProm(defaultPromtailIP,port,promtailPath,request,response);
     	
     }
     
     public void requestProm(String host, int port, String path, HttpServletRequest request, HttpServletResponse response)  {
         String queryString = ReflectUtil.invoke(request, "getQueryString");
         queryString = StrUtil.isNotBlank(queryString)?queryString:"";
         String url = UrlBuilder.create().setScheme("http").setHost(host).setPort(port).setPath(UrlPath.of(path, Charset.forName("UTF-8"))).toURL().toString() + "?" + queryString;

         log.info("promtail url: {}", url);
         String method = request.getMethod();
         HttpURLConnection conn = null;
         ServletInputStream reqInputStream = null;
         ServletOutputStream resOutputStream = null;
         OutputStream connOutputStream = null;
         InputStream connInputStream = null;
         try {
             conn = HttpConnection.create(url, null).getHttpURLConnection();
             reqInputStream = request.getInputStream();
             resOutputStream = response.getOutputStream();
             conn.setRequestMethod(method);
             // 复制请求头
             Enumeration<String> headerNames = request.getHeaderNames();
             while (headerNames.hasMoreElements()) {
                 String hn = headerNames.nextElement();
                 if(!"authorization".equalsIgnoreCase(hn)) {
                 	ReflectUtil.invoke(conn,"addRequestProperty",hn,request.getHeader(hn));
                 }
             }
             
             if (!"GET".equalsIgnoreCase(method)) {
                 conn.setDoOutput(true);
                 conn.setDoInput(true);
                 conn.setUseCaches(false);
                 connOutputStream = conn.getOutputStream();
                 Tool.IoUtil.copy(reqInputStream, connOutputStream);
             }
             conn.connect();
             int responseCode = conn.getResponseCode();
             connInputStream = (responseCode < 400)? conn.getInputStream():conn.getErrorStream();
             String responseMessage = conn.getResponseMessage();
             Map<String, List<String>> responseHeaders = conn.getHeaderFields();
             //复制响应头
             for(Map.Entry<String, List<String>> en : responseHeaders.entrySet()) {
                 String key = en.getKey();
                 if (Tool.StrUtil.isEmpty(key) || "Transfer-Encoding".equals(key)) continue;
                 List<String> value = en.getValue();
                 ReflectUtil.invoke(response,"addHeader",key,Tool.StrUtil.join("; ",value));
             }
             ReflectUtil.invoke(response, "setStatus", responseCode, responseMessage);

             Tool.IoUtil.copy(connInputStream, resOutputStream);
             resOutputStream.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(reqInputStream,resOutputStream,connOutputStream,connInputStream);
             if(conn != null){
                 conn.disconnect();
             }
         }
     }

}