mini-chat/cmd/serverFunc.go

/*
Copyright © 2024 Raul <raul@bulgariu.xyz>
*/

package cmd

import (
	"bufio"
	"crypto/tls"
	_ "embed"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"strings"
)

var (
	listenPort   string = "1302"
	password     string = ""
	isLogging    bool   = false
	logLocation  string
	listenerList []chan string
	servInsecure bool
)

//go:embed gen-cert.sh
var script string

type User struct {
	Username string
	IP       string
}

func (u User) CreateUser(usr string, ip string) User {
	u.Username = usr
	u.IP = ip
	return u
}

func createCerts() {
	fmt.Println("[-] Certificates don't exist! Creating them...")
	c := exec.Command("bash")
	c.Stdin = strings.NewReader(script)
	b, err := c.Output()
	if err != nil {
		log.Fatalf("Error occurred creating certificates: %v\n", err)
	}
	fmt.Print(string(b))
}

func startInsecureServer() (net.Listener, error) {
	ln, err := net.Listen("tcp", ":"+listenPort)
	return ln, err
}

func startSecureServer() (net.Listener, error) {
	cer, err := tls.LoadX509KeyPair("server.crt", "server.key")
	if os.IsNotExist(err) {
		createCerts()
		cer, err = tls.LoadX509KeyPair("server.crt", "server.key")
	}
	if err != nil {
		log.Fatalf("Error happened loading certificates: %v\n", err)
	}
	config := &tls.Config{Certificates: []tls.Certificate{cer}}
	ln, err := tls.Listen("tcp", ":"+listenPort, config)
	return ln, err
}

func Server() {
	var ln net.Listener
	var err error
	if servInsecure == true {
		fmt.Println("WARNING: Starting unencrypted server!")
		ln, err = startInsecureServer()
	} else {
		ln, err = startSecureServer()
	}

	if err != nil {
		log.Fatalf("Error happened trying to listen on port: %v\n", err)
	}
	defer ln.Close()
	fmt.Printf("Listening on port %v...\n", listenPort)
	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Fatalf("Error happened trying to accept connection: %v\n", err)
		}
		chatChan := make(chan string, 30)
		listenerList = append(listenerList, chatChan)
		go handleConn(conn, chatChan)
	}
}

func getUsername(conn net.Conn) (s string, err error) {
	conn.Write([]byte("What's your name?\nChoice: "))
	name, err := bufio.NewReader(conn).ReadString('\n')
	if err != nil {
		return "", err
	}
	trimmedName := strings.TrimRight(name, "\n")
	return trimmedName, nil
}

func getUserInput(conn net.Conn) (s string, err error) {
	message, err := bufio.NewReader(conn).ReadString('\n')
	if err != nil {
		return "", err
	}

	return message, nil
}

func removeFromList(chatChan chan string) {
	for i, v := range listenerList {
		if v == chatChan {
			listenerList = append(listenerList[:i], listenerList[:i+1]...)
		}
	}
}

func populateChat(conn net.Conn) {
	if isLogging == false {
		return
	}
	file, err := os.Open(logLocation)
	if err != nil {
		log.Printf("Error opening file for populating: %v\n", err)
	}
	defer file.Close()
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		conn.Write([]byte(fmt.Sprintln(scanner.Text())))
	}

}

func getPasswd(conn net.Conn) error {
	conn.Write([]byte("Password: "))
	userPassNewline, err := bufio.NewReader(conn).ReadString('\n')
	userPass := strings.TrimRight(userPassNewline, "\n")
	if err != nil {
		e := fmt.Errorf("Node %v didn't respond to password prompt!\n", getIP(conn))
		return e
	}
	if userPass != password {
		e := fmt.Errorf("Node %v attempted connecting with an incorrect password!\n", getIP(conn))
		return e
	}
	return nil
}

func handleConn(conn net.Conn, chatChan chan string) {
	defer conn.Close()

	if password != "" {
		if err := getPasswd(conn); err != nil {
			log.Print(err)
			return
		}
	}

	go receiveMessageServer(conn, chatChan)

	//////////////////////////////////
	// Get user information
	//////////////////////////////////
	userName, err := getUsername(conn)
	if err != nil {
		log.Printf("Error occurred getting username: %v\n", err)
		//removeFromList(chatChan)
		return
	}
	userIP := getIP(conn)
	//////////////////////////////////
	populateChat(conn)

	newUserTemplate := new(User)
	newUser := newUserTemplate.CreateUser(userName, userIP)
	joinMessage := fmt.Sprintf("%v has joined the chat!", newUser.Username)
	fmt.Println(joinMessage)
	addToLog(fmt.Sprintln(joinMessage))
	//conn.Write([]byte(joinMessage))
	sendMessage(joinMessage)

	//////////////////////////////////
	for {
		message, err := getUserInput(conn)
		if err != nil {
			quitMessage := fmt.Sprintf("%v has disconnected!", newUser.Username)
			fmt.Println(quitMessage)
			addToLog(fmt.Sprintln(quitMessage))
			sendMessage(quitMessage)
			//removeFromList(chatChan)
			// if _, err := conn.Write([]byte(quitMessage)); err != nil {
			// 	log.Printf("Error happened sending disconnect message: %v", err)
			// }
			return
		}
		finalMessage := fmt.Sprintf("[%v] %v: %v", newUser.IP, newUser.Username, strings.TrimRight(message, "\n"))
		fm := fmt.Sprintf("%v\n", finalMessage)
		fmt.Print(fm)
		addToLog(fm)
		sendMessage(finalMessage)

		//chatChan <- finalMessage
		// if _, err := conn.Write([]byte(finalMessage)); err != nil {
		// 	log.Printf("Error happened sending message: %v", err)
		// }
	}
	//////////////////////////////////
}

func sendMessage(msg string) {
	for _, ch := range listenerList {
		ch <- msg
	}
}

func receiveMessageServer(conn net.Conn, chatChan chan string) {
	for {
		select {
		case message := <-chatChan:
			conn.Write([]byte(message))
		}
	}
}

func getIP(conn net.Conn) (IP string) {
	if addr, ok := conn.RemoteAddr().(*net.TCPAddr); ok {
		IP = fmt.Sprintf("%v", addr.IP)
	}
	return IP
}

func addToLog(s string) {
	if isLogging == false {
		return
	}
	file, err := os.OpenFile(logLocation, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0640)
	if err != nil {
		log.Printf("Error occurred: %v\n", err)
	}
	defer file.Close()
	file.WriteString(s)
}