package main

import (
	"database/sql"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"golang.org/x/crypto/ssh"
	"gopkg.in/ini.v1"
)

const (
	DAEMON      = "d"
	FOREVER     = "f"
	HOST_IP     = "h"
	HOST_PORT   = "p"
	HOST_USER   = "u"
	HOST_REMOTE = "r"
	HOST_LOCAL  = "l"
)

func StripSlice(slice []string, element string) []string {
	for i := 0; i < len(slice); {
		if slice[i] == element && i != len(slice)-1 {
			slice = append(slice[:i], slice[i+1:]...)
		} else if slice[i] == element && i == len(slice)-1 {
			slice = slice[:i]
		} else {
			i++
		}
	}
	return slice
}

func SubProcess(args []string) *exec.Cmd {
	cmd := exec.Command(args[0], args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Start()
	if err != nil {
		fmt.Fprintf(os.Stderr, "[-] Error: %s\n", err)
	}
	return cmd
}

// Check if a port is available
func Check(port int) (status bool, err error) {

	// Concatenate a colon and the port
	host := ":" + strconv.Itoa(port)

	// Try to create a server with the port
	server, err := net.Listen("tcp", host)

	// if it fails then the port is likely taken
	if err != nil {
		return false, err
	}

	// close the server
	server.Close()

	// we successfully used and closed the port
	// so it's now available to be used again
	return true, nil

}

func GetCurrentDirectory() string {
	dir, err := filepath.Abs(filepath.Dir(os.Args[0])) // 返回绝对路径  filepath.Dir(os.Args[0])去除最后一个元素的路径
	if err != nil {
		log.Fatal(err)
	}
	return strings.Replace(dir, "\\", "/", -1) // 将\替换成/
}

func PathExists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}

	if os.IsNotExist(err) {
		return false, nil
	}

	return false, err
}

type user struct {
	id               string
	host_id          string
	host_ip          string
	router_ip        string
	router_port      string
	protocol_type    string
	protocol_port    string
	state            string
	acc_auth_id      string
	auth_type        string
	username         string
	username_prompt  string
	password_prompt  string
	password         string
	password_decrypt string
	pri_key          string
	creator_id       string
	create_time      string
	last_secret      string
}

func ExecCommand(strCommand string) string {
	cmd := exec.Command("/bin/bash", "-c", strCommand)

	stdout, _ := cmd.StdoutPipe()
	if err := cmd.Start(); err != nil {
		fmt.Println("Execute failed when Start: " + err.Error())
		return ""
	}

	out_bytes, _ := ioutil.ReadAll(stdout)
	stdout.Close()

	if err := cmd.Wait(); err != nil {
		fmt.Println("Execute failed when Wait: " + err.Error())
		return ""
	}

	return string(out_bytes)
}

// 读取mysql数据库内信息
func mysql_show(connect_info string, host_ip string, host_user string, host_port string) user {

	var db *sql.DB
	db, _ = sql.Open("mysql", connect_info) // 设置连接数据库的参数
	defer db.Close()                        // 关闭数据库
	err := db.Ping()                        // 连接数据库
	if err != nil {
		// 退出
		log.Printf("数据库连接失败")
		syscall.Kill(os.Getppid(), syscall.SIGKILL)
		syscall.Kill(os.Getpid(), syscall.SIGKILL)
	} else {
		log.Printf("数据库连接成功")

	}
	sqlStr := "select host_ip, username, protocol_port, password from tp_acc where host_ip = ? and username = ? and protocol_port = ?"
	rows, err := db.Query(sqlStr, host_ip, host_user, host_port)
	if err != nil {
		log.Printf("query failed, err:%v\n", err)
	}
	// 非常重要关闭rows释放持有的数据库链接
	defer rows.Close()
	var u user
	// 循环读取结果集中的数据
	for rows.Next() {
		err := rows.Scan(&u.host_ip, &u.username, &u.protocol_port, &u.password)
		if err != nil {
			log.Printf("scan failed, err:%v\n", err)
		}
		log.Printf("%s %s %s %s\n", u.host_ip, u.protocol_port, u.username, u.password)
	}

	// 解密密钥
	command := "./tp_decrypt "
	command = command + u.password
	strData := ExecCommand(command)
	strData = strings.Replace(strData, "\n", "", -1)
	fmt.Println(strData)
	if strData == "" {
		// 判断密码是否解密成功，不成功退出
		syscall.Kill(os.Getppid(), syscall.SIGKILL)
		syscall.Kill(os.Getpid(), syscall.SIGKILL)
	}
	u.password_decrypt = strData

	return u
}

func main() {
	// 参数处理
	daemon := flag.Bool(DAEMON, false, "run in daemon")
	forever := flag.Bool(FOREVER, false, "run forever")
	host_ip := flag.String(HOST_IP, "", "DB服务器SSH IP地址")
	host_port := flag.String(HOST_PORT, "", "DB服务器SSH PORT")
	host_user := flag.String(HOST_USER, "", "DB服务器SSH USER用户")
	remote_port := flag.String(HOST_REMOTE, "", "DB服务器端口（如: 3306、1521 ...）")
	local_port := flag.String(HOST_LOCAL, "", "本地监听端口（或者堡垒机监听端口）")
	connect_max := flag.Int("c", 999, "建立隧道最大次数（与DB服务器断开）")
	flag.Parse()

	// 判断配置文件是否存在
	INIFILE := GetCurrentDirectory() + "/" + "tunnel.ini"
	b, _ := PathExists(INIFILE)
	if !b {
		INIFILE = "/etc/tunnel.ini"
	}
	// 读取配置文件
	cfg, inierr := ini.Load(INIFILE)
	if inierr != nil {
		fmt.Printf("Fail to read file: %v", inierr)
		os.Exit(1)
	}
	// 读取数据库连接信息
	MYSQL_INFO := cfg.Section("global").Key("MYSQL_INFO").String()
	if MYSQL_INFO == "" {
		log.Printf("读取Mysql配置出错!\n")
		os.Exit(0)
	}

	if *daemon {
		SubProcess(StripSlice(os.Args, "-"+DAEMON))
		fmt.Printf("[*] Daemon running in PID: %d PPID: %d\n", os.Getpid(), os.Getppid())
		os.Exit(0)
	} else if *forever {
		for {
			cmd := SubProcess(StripSlice(os.Args, "-"+FOREVER))
			fmt.Printf("[*] Forever running in PID: %d PPID: %d\n", os.Getpid(), os.Getppid())
			time.Sleep(time.Second * 5)
			cmd.Wait()
		}
		//os.Exit(0)
	} else {
		fmt.Printf("[*] Service running in PID: %d PPID: %d\n", os.Getpid(), os.Getppid())
	}

	local_port_, _ := strconv.Atoi(*local_port)
	r, _ := Check(local_port_)
	if r {
		u := mysql_show(MYSQL_INFO, *host_ip, *host_user, *host_port)
		tunnel := NewSSHTunnel(u.username+"@"+u.host_ip+":"+u.protocol_port, ssh.Password(u.password_decrypt), "0.0.0.0:"+*remote_port, *local_port)
		tunnel.Log = log.New(os.Stdout, "", log.Ldate|log.Lmicroseconds)

		tunnel.Start(*connect_max)
		tunnel.Close()
		time.Sleep(100 * time.Millisecond)
	} else {
		log.Printf("端口: %d 不可用！退出！\n", local_port_)
		log.Printf("DoSomething running in PPID: %d\n", os.Getppid())

		syscall.Kill(os.Getppid(), syscall.SIGKILL)
		time.Sleep(time.Second * 5)
		os.Exit(0)
	}

	time.Sleep(time.Second * 5)
	os.Exit(0)
}
