CSI-ES-2324/Projs/PD1/internal/client/client.go

package client

import (
	"PD1/internal/protocol"
	"PD1/internal/utils/cryptoUtils"
	"PD1/internal/utils/networking"
	"crypto/x509"
	"flag"
	"fmt"
	"sort"
)

func Run() {
	var userFile string
	flag.StringVar(&userFile, "user", "userdata.p12", "Specify user data file")
	flag.Parse()

	if flag.NArg() == 0 {
		panic("No command provided. Use 'help' for instructions.")
	}
	//Get user KeyStore
	password := AskUserPassword()
	clientKeyStore := cryptoUtils.LoadKeyStore(userFile, password)

	command := flag.Arg(0)
	switch command {
	case "send":
		if flag.NArg() < 3 {
			panic("Insufficient arguments for 'send' command. Usage: send <UID> <SUBJECT>")
		}
		uid := flag.Arg(1)
		subject := flag.Arg(2)
		messageBody := readMessageBody()
		//Turn content to bytes
		marshaledSubject := Marshal(subject)
		marshaledBody := Marshal(messageBody)

		cl := networking.NewClient[protocol.Packet](&clientKeyStore)
		defer cl.Connection.Conn.Close()

		uidCert := getUserCert(cl, uid)
		if uidCert == nil {
			return
		}
		encryptedSubject := clientKeyStore.EncryptMessageContent(uidCert, marshaledSubject)
		encryptedBody := clientKeyStore.EncryptMessageContent(uidCert, marshaledBody)
		submitMessage := protocol.NewSubmitMessagePacket(uid, encryptedSubject, encryptedBody)
		if !cl.Connection.Send(submitMessage) {
			return
		}
		cl.Connection.Conn.Close()

	case "askqueue":
		cl := networking.NewClient[protocol.Packet](&clientKeyStore)
		defer cl.Connection.Conn.Close()

		requestUnreadMsgsQueuePacket := protocol.NewRequestUnreadMsgsQueuePacket()
		if !cl.Connection.Send(requestUnreadMsgsQueuePacket) {
			return
		}
		serverMessagePackets, certificates := getManyMessagesInfo(cl)
		var clientMessages []ClientMessageInfo
		for _, message := range serverMessagePackets {
			senderCert, ok := certificates[message.FromUID]
			if ok {
				decryptedSubjectBytes := clientKeyStore.DecryptMessageContent(senderCert, message.Subject)
				subject := Unmarshal(decryptedSubjectBytes)
				clientMessage := newClientMessageInfo(message.Num, message.FromUID, subject, message.Timestamp)
				clientMessages = append(clientMessages, clientMessage)
			}
		}
		//Sort the messages
		sort.Slice(clientMessages, func(i, j int) bool {
			return clientMessages[i].Num > clientMessages[j].Num
		})

		showMessagesInfo(clientMessages)

	//case "getall":
	//	cl := networking.NewClient[protocol.Packet](&clientKeyStore)
	//	defer cl.Connection.Conn.Close()

	//	requestAllMsgPacket := protocol.NewRequestAllMsgPacket()
	//	if !cl.Connection.Send(requestAllMsgPacket) {
	//		return
	//	}
	//    serverMessagePackets,certificates := getManyMessages(cl)
	//	var clientMessages []ClientMessage
	//	for _, message := range serverMessagePackets {
	//		senderCert, ok := certificates[message.FromUID]
	//		if ok {
	//			decryptedContentBytes := clientKeyStore.DecryptMessageContent(senderCert, message.Content)
	//			content := UnmarshalContent(decryptedContentBytes)
	//			clientMessage := newClientMessage(message.FromUID, message.ToUID, content, message.Timestamp)
	//			clientMessages = append(clientMessages, clientMessage)
	//		}
	//	}
	//	//Sort the messages
	//	sort.Slice(clientMessages, func(i, j int) bool {
	//		return clientMessages[i].Timestamp.After(clientMessages[j].Timestamp)
	//	})

	//	showMessages(clientMessages)

	case "getmsg":
		if flag.NArg() < 2 {
			panic("Insufficient arguments for 'getmsg' command. Usage: getmsg <NUM>")
		}
		//num := flag.Arg(1)
		cl := networking.NewClient[protocol.Packet](&clientKeyStore)
		defer cl.Connection.Conn.Close()

	case "help":
		showHelp()

	default:
		commandError()
	}

}

func getUserCert(cl networking.Client[protocol.Packet], uid string) *x509.Certificate {
	certRequestPacket := protocol.NewRequestUserCertPacket(uid)
	if !cl.Connection.Send(certRequestPacket) {
		return nil
	}
	var certPacket *protocol.Packet
	certPacket, active := cl.Connection.Receive()
	if !active {
		return nil
	}
	uidCertInBytes := protocol.UnmarshalSendUserCertPacket(certPacket.Body)
	uidCert, err := x509.ParseCertificate(uidCertInBytes.Certificate)
	if err != nil {
		return nil
	}
	return uidCert
}

func getManyMessagesInfo(cl networking.Client[protocol.Packet]) ([]protocol.ServerMessageInfoPacket, map[string]*x509.Certificate) {
	//Create the slice to hold the incoming messages before decrypting
	//Create the map to hold the sender certificates
	//Create sync mutexes
	serverMessageInfoPackets := []protocol.ServerMessageInfoPacket{}
	//Run while message isn't the last one
    msg := protocol.ServerMessageInfoPacket{}
	for !msg.Last {
		sendMsgPacket, active := cl.Connection.Receive()
		if !active {
			return nil, nil
		}
		msg = protocol.UnmarshalServerMessageInfoPacket(sendMsgPacket.Body)
		//Lock and append
		serverMessageInfoPackets = append(serverMessageInfoPackets, msg)
	}

	//Create Set of needed certificates
	senderSet := map[string]bool{}
	for _, messageInfo := range serverMessageInfoPackets {
		senderSet[messageInfo.FromUID] = true
	}
	certificatesMap := map[string]*x509.Certificate{}
	//Get senders' certificates
	for senderUID := range senderSet {
		senderCert := getUserCert(cl, senderUID)
        fmt.Println("Got a User cert")
		certificatesMap[senderUID] = senderCert
	}
	return serverMessageInfoPackets, certificatesMap
}