package main

import (
	"crypto/tls"
	"log"
	"net"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"syscall"
)

var (
	cfile = flag.String("cert", "cert.pem", "Certificate file in PEM format")
	kfile = flag.String("key", "key.pem", "Key file in PEM format")
	port = flag.Int("port", 1234, "Port to bind to")
	uid = flag.Int("uid", -1, "UID to run under")
	gid = flag.Int("gid", -1, "GID to run under")
	args []string
	nargs int
)

func main() {

	flag.Parse()
	 args = flag.Args()
	nargs = flag.NArg()
	if nargs < 1 {
		fmt.Println("Usage: tlsserver [options] cmd [flags for cmd]")
		fmt.Println("options:")
		flag.PrintDefaults()
		os.Exit(1)
	}

	// setup certs etc. for TLS-socket
	tconf := new(tls.Config)
	cert, err := tls.LoadX509KeyPair(*cfile, *kfile)
	if err != nil {
		fmt.Println("error with certs:", err)
		os.Exit(2)
	}

	tconf.Certificates = append(tconf.Certificates, cert)
	tconf.BuildNameToCertificate()

	// start listening
	sport := fmt.Sprintf(":%d", *port)
	sock, err := tls.Listen("tcp", sport , tconf)
	if err != nil {
		fmt.Println("error with tcp-socket:", err)
		os.Exit(3)
	}
	defer sock.Close()

	// set uid/gid
	if *gid >= 0 {
		err := syscall.Setgid(*gid)
		if err != nil {
			fmt.Println("Couldn't setgid to", *gid, ":", err)
			os.Exit(4)
		}
	}

	if *uid >= 0 {
		err := syscall.Setuid(*uid)
		if err != nil {
			fmt.Println("Couldn't setuid to", *uid, ":", err)
			os.Exit(4)
		}
	}



	// accept-loop
	for {
		conn, err := sock.Accept()
		if err != nil {
			log.Println("error during Accept()", err)
			continue
		}
		log.Println("Got connection:", conn.RemoteAddr())
		go handleConnection(conn)
	}
}

func handleConnection(conn net.Conn) {
	defer conn.Close()

	// setup cmd
	cmd := exec.Command(args[0])
	cmd.Args = args
	cmd.Stdin = conn
	cmd.Stdout = conn
	cmd.Stderr = os.Stderr

	// prepare environment according to tcp-environ(5)
	lh, lp, err := net.SplitHostPort(conn.LocalAddr().String())
	if err != nil {
		log.Println(err)
		return
	}
	rh, rp, err := net.SplitHostPort(conn.LocalAddr().String())
	if err != nil {
		log.Println(err)
		return
	}
	cmd.Env = make([]string,0)
	cmd.Env = append(cmd.Env, "PATH="+os.Getenv("PATH"))
	cmd.Env = append(cmd.Env, "PROTO=TCP")
	cmd.Env = append(cmd.Env, "TCPLOCALIP="+lh)
	cmd.Env = append(cmd.Env, "TCPLOCALPORT="+lp)
	cmd.Env = append(cmd.Env, "TCPREMOTEIP="+rh)
	cmd.Env = append(cmd.Env, "TCPREMOTEPORT="+rp)

	err = cmd.Run()
	if err != nil {
		log.Println("after Run: ", err)
	}
	log.Println("Done with connection", conn.RemoteAddr())
}