package main

import (
	"archive/tar"
	"compress/gzip"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"crypto/tls"
	"os"
	"path/filepath"
	"strings"
	"time"
)

type Session struct {
	ID        string `json:"id"`
	DateStart string `json:"date_start"`
	Protocol  string `json:"protocol"`
}

func main() {
	//忽略SSL证书告警
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	
	client := &http.Client{Transport: tr}
	
	// 获取命令行参数
	url := flag.String("url", "", "URL")
	flag.Parse()

	if *url == "" {
		log.Fatal("请传入有效的URL")
	}

	// 发送请求获取JSON数据
	resp, err := client.Get(*url + "/api/v1/terminal/sessions/?limit=1")
	if err != nil {
		log.Fatal("请求失败:", err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("读取响应失败:", err)
	}

	// 解析JSON数据
	var sessionData map[string]interface{}
	err = json.Unmarshal(body, &sessionData)
	if err != nil {
		log.Fatal("解析JSON失败:", err)
	}

	// 获取session信息
	results := sessionData["results"].([]interface{})
	if len(results) == 0 {
		log.Fatal("未找到session信息")
	}

	session := results[0].(map[string]interface{})
	sessionID := session["id"].(string)
	dateStart := session["date_start"].(string)
	protocol := session["protocol"].(string)

	// 将时间格式转换为YYYY-MM-DD
	parsedTime, err := time.Parse("2006/01/02 15:04:05 -0700", dateStart)
	if err != nil {
		log.Fatal("时间转换失败:", err)
	}
	dateStartFormatted := parsedTime.Format("2006-01-02")

	// 保存返回的JSON数据
	jsonBytes, err := json.Marshal(session)
	if err != nil {
		log.Fatal("序列化JSON失败:", err)
	}

	err = ioutil.WriteFile(sessionID+".json", jsonBytes, 0644)
	if err != nil {
		log.Fatal("保存JSON文件失败:", err)
	}

	// 根据协议定义replay_url
	var replayURL string
	if strings.ToLower(protocol) == "rdp" {
		replayURL = ".replay.gz"
	} else {
		replayURL = ".cast.gz"
	}

	// 发送请求获取replay内容并保存
	replayURL1 := "/media/xpack/../replay/" + dateStartFormatted + "/" + sessionID + replayURL
	resp, err = client.Get(*url + replayURL1)
	if err != nil {
		log.Fatal("请求replay失败:", err)
	}
	defer resp.Body.Close()

	cacheBytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("读取replay内容失败:", err)
	}
	
	// 创建输出文件
	file, err := os.Create(sessionID + replayURL)
	if err != nil {
		fmt.Println("无法创建输出文件:", err)
		return
	}
	defer file.Close()

	// 创建gzip.Writer，并设置文件名
	gzipWriter := gzip.NewWriter(file)
	gzipWriter.Header.Name = sessionID

	// 写入需要压缩的数据
	_, err = gzipWriter.Write([]byte(cacheBytes))
	if err != nil {
		fmt.Println("写入压缩数据时发生错误:", err)
		return
	}

	// 关闭gzip.Writer，在关闭之前会将压缩数据写入到输出文件
	err = gzipWriter.Close()
	if err != nil {
		fmt.Println("关闭gzip.Writer时发生错误:", err)
		return
	}

	// 创建tar文件
	tarFile, err := os.Create(sessionID + ".tar")
	if err != nil {
		log.Fatal("创建tar文件失败:", err)
	}
	defer tarFile.Close()

	tarWriter := tar.NewWriter(tarFile)
	defer tarWriter.Close()

	// 添加sessionID.json到tar文件
	err = addFileToTar(sessionID+".json", tarWriter)
	if err != nil {
		log.Fatal("添加文件到tar失败:", err)
	}

	// 添加sessionID.replay.gz到tar文件
	err = addFileToTar(sessionID+replayURL, tarWriter)
	if err != nil {
		log.Fatal("添加文件到tar失败:", err)
	}

	fmt.Println("脚本执行成功！")
}

// 添加文件到tar
func addFileToTar(filePath string, tarWriter *tar.Writer) error {
	file, err := os.Open(filePath)
	if err != nil {
		return err
	}
	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		return err
	}

	header := &tar.Header{
		Name:    filepath.Base(filePath),
		Mode:    int64(stat.Mode()),
		Size:    stat.Size(),
		ModTime: stat.ModTime(),
	}

	err = tarWriter.WriteHeader(header)
	if err != nil {
		return err
	}

	_, err = io.Copy(tarWriter, file)
	if err != nil {
		return err
	}

	return nil
}