package main

import (
    "encoding/json"
    "flag"
    "fmt"
    "io"
    "log"
    "net"
    "net/http"
    "os"
    "os/signal"
    "runtime"
    "sync"
    "syscall"
    "time"
)

var restApiServer = flag.String("restApi", "", "listen addr for restapi") //监听rest full api 地址   ip+端口 默认值 0.0.0.0:8000  参见参数判断
var auth = flag.String("auth", "taven123", "restApi Password")  //身份认证   默认rest api  是  taven123
var gLocalConn net.Listener   //监听地址

var clientSMap map[string]net.Conn

var forwardInfo string

func main() {

    clientSMap = make(map[string]net.Conn) //key 为 ip   value 为  ip+端口创建的链接

    //解析传入的参数
    flag.Parse()

    if *restApiServer == "" {   //默认监听地址设置
        *restApiServer = "0.0.0.0:8000"  //默认值
    }

    go StartHttpServer(*restApiServer)//开启http 监听

    log.Println("restApiServer:", *restApiServer)
    fmt.Println("------------启动成功------------")

    //开启线程同步锁
    var w sync.WaitGroup
    w.Add(2)  //默认监听两个链接  

    //开一个并发线程,接收退出信号
    go func() {
        c := make(chan os.Signal, 1)  //信号量通道
        signal.Notify(c, os.Interrupt, syscall.SIGTERM)  //监听系统信号  
        n := 0
        f := func() {
            <-c
            n++
            if n > 2 {
                log.Println("force shutdown")
                os.Exit(-1)
            }
            log.Println("received signal,shutdown")
            closeAllConn() //关闭对应客户端ip 对应的链接
        }
        f()
        go func() {
            for {
                f()
            }
        }()
        //执行完成一次,Done() 等同于 Add(-1),计数不为0,则阻塞
        w.Done()
    }()

    loop := func() {
        w.Done()

    }
    loop()
    w.Wait()

    fmt.Println("------------程序执行完成------------")

}

func StartHttpServer(addr string) {  //开启服务器端的信息监听

    http.HandleFunc("/ServerSummary", ServerSummary)  //获取服务状态
    http.HandleFunc("/ForwardWork", ForwardWork)  //端口映射服务

    //
    err := http.ListenAndServe(addr, http.DefaultServeMux)

    if err != nil {
        fmt.Println("ListenAndServe error: ", err.Error())
    }

}

func ServerSummary(rw http.ResponseWriter, req *http.Request) {
    log.Println("ServerSummary")
    obj := make(map[string]interface{})
    obj["runtime_NumGoroutine"] = runtime.NumGoroutine()  //当前开启协程数量
    obj["runtime_GOOS"] = runtime.GOOS  //当前操作系统类型
    obj["runtime_GOARCH"] = runtime.GOARCH //操作系统架构
    obj["restApi_Addr"] = *restApiServer   //请求rest服务
    obj["server_Time"] = time.Now()
    obj["clients_Count"] = len(clientSMap)   //客户端数量

    var clist []string
    for cId, _ := range clientSMap {  //迭代客户端来链接map
        clist = append(clist, cId)
    }
    obj["clients_List"] = clist   //客户端列表
    obj["forwardInfo"] = forwardInfo  //转向服务器信息

    res, err := json.Marshal(obj)  //生成json数据
    if err != nil {
        log.Println("json marshal:", err)
        return
    }

    rw.Header().Add("Content-Type", "application/json;charset=utf-8")  //添加响应数据头 数据格式
    _, err = rw.Write(res)
    if err != nil {
        log.Println("write err:", err)
    }
    return
}

func ForwardWork(rw http.ResponseWriter, req *http.Request) {  
    req.ParseForm()  //解析请求参数

    obj := make(map[string]interface{}) //存储状态信息
    obj["code"] = 0
    obj["msg"] = ""

    paramAuth, hasAuth := req.Form["auth"]  //获取参数  密码
    if !hasAuth {   //情况一   不存在密码
        log.Println("request no auth")
        obj["code"] = 1
        obj["msg"] = "request no auth"
        responseResult(obj, rw)  //json数据格式 响应数据
        return

    }

    if paramAuth[0] != *auth {   //情况二  存在密码    但是密码不匹配
        log.Println("request auth failed")
        obj["code"] = 1
        obj["msg"] = "request auth failed"
        responseResult(obj, rw)//json数据格式 响应数据

        return
    }

    paramStatus, hasStatus := req.Form["status"]  //status:如果是开启转发,则为1,如果是关闭转发,则为0
    if !hasStatus {  //没有开启转发  直接返回
        return

    }

    log.Println("param_status:", paramStatus)

    if paramStatus[0] == "1" {  //开启转发
        //启动服务
        paramFromAddr, hasFromAddr := req.Form["fromAddr"]  //要用来在A机器上监听的一个端口,用来给客户端连接     =====服务器
        paramToAddr, hasToAddr := req.Form["toAddr"] //把fromAddr端口的数据转发到哪个IP的端口上   ===转发到的服务器
        if gLocalConn != nil { //目标服务器是否存在链接   不存在 直接拒接
            gLocalConn.Close()  //释放资源
        }

        if hasFromAddr && hasToAddr {  //目标服务器 和 转发到服务器存在   
            go forwardPort(paramFromAddr[0], paramToAddr[0])
        }
    }

    if paramStatus[0] == "0" {   //转发状态为0 时   直接释放资源 
        //关闭服务
        closeAllConn()
        forwardInfo = ""
    }

    responseResult(obj, rw)// json格式数据响应客户端

    return

}

func responseResult(data map[string]interface{}, rw http.ResponseWriter) {// json格式数据响应客户端
    res, err := json.Marshal(data)  //生成json数据格式
    if err != nil {
        log.Println("json marshal:", err)
        return
    }

    rw.Header().Add("Content-Type", "application/json;charset=utf-8") //响应头数据格式设置
    _, err = rw.Write(res) //响应写入数据
    if err != nil {
        log.Println("write err:", err)
    }
}

func closeAllConn() {// 关闭连接
    for cId, conn := range clientSMap {
        log.Println("clientMap id:", cId)
        conn.Close() //关闭连接 conn
        delete(clientSMap, cId)  //从客户端map中删除对应的数据
    }

    if gLocalConn != nil {  //关闭元数据链接【服务端】
        gLocalConn.Close()
        log.Println("Listener Close")
    } else {
        gLocalConn = nil
        log.Println("Listener set to nil", gLocalConn)
    }
}

func forwardPort(sourcePort string, targetPort string) { //转发服务

    fmt.Println("sourcePort:", sourcePort, "targetPort:", targetPort)

    localConn, err := net.Listen("tcp", sourcePort) //目标服务器   监听服务器

    if err != nil {
        fmt.Println(err.Error())
        return
    }

    gLocalConn = localConn

    fmt.Println("服务启动成功,服务地址:", sourcePort)

    forwardInfo = fmt.Sprintf("%s - %s", sourcePort, targetPort)

    for {
        fmt.Println("Ready to Accept ...")
        sourceConn, err := gLocalConn.Accept()  //获取监听服务链接

        if err != nil {
            log.Println("server err:", err.Error())
            break
        }
        //log.Println("client", sc.id, "create session", sessionId)

        id := sourceConn.RemoteAddr().String()   //获取客户端ip地址
        clientSMap[id] = sourceConn

        fmt.Println("conn.RemoteAddr().String() :", id)

        //targetPort := "172.16.128.83:22"
        targetConn, err := net.DialTimeout("tcp", targetPort, 30*time.Second)//获取转发服务链接
               //开启协程 支持从客户端 到服务端  以及 从 服务端到客户端数据 读写
        go func() {
            _, err = io.Copy(targetConn, sourceConn)  //从sourceConn读取数据 写入到targetConn中
            if err != nil {
                //log.Fatalf("io.Copy 1 failed: %v", err)
                fmt.Println("io.Copy 1 failed:", err.Error())
            }
        }()

        go func() {
            _, err = io.Copy(sourceConn, targetConn)//从targetConn读取数据 写入到sourceConn中
            if err != nil {
                //log.Fatalf("io.Copy 2 failed: %v", err)
                fmt.Println("io.Copy 2 failed:", err.Error())
            }
        }()

    }

    //
    log.Println("forwardPort end.")

}

05-11 15:19