DHT-kademlia-P2Psec/commandlineDHT/main.go

package main

import (
	"encoding/binary"
	"errors"
	"flag"
	kademlia "gitlab.lrz.de/netintum/teaching/p2psec_projects_2024/DHT-6/pkg"
	"io"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
)

// Config holds the config variables for the DHT module

func main() {

	mode := flag.String("m", "", "Whether to send a put or get packet")
	ttl := flag.Int("ttl", 3600, "TTL for sending a put or get packet")
	replication := flag.Int("rep", 20, "Replication factor for put packet")
	address := flag.String("ip", "", "IP and port of API Server in for 127.0.0.1:1234")
	keyStr := flag.String("key", "", "Key to be stored or retrieved")
	valueStr := flag.String("value", "", "Value to pe put")
	certPath := flag.String("cert", "", "Path to certificate")
	flag.Parse()

	var key [32]byte
	copy(key[:], *keyStr)
	value := []byte(*valueStr)

	switch *mode {
	case "put":
		conn, err := net.Dial("tcp", *address)
		if err != nil {
			log.Fatal(err)
		}
		packet := kademlia.DHTPutPacket{
			TTL:         uint16(*ttl),
			Replication: byte(*replication),
			Key:         key,
			Value:       value,
		}
		conn.Write(kademlia.SerialisePutPacket(&packet))
		conn.Close()
	case "get":
		conn, err := net.Dial("tcp", *address)
		if err != nil {
			log.Fatal(err)
		}
		packet := kademlia.DHTGetPacket{Key: key}
		conn.Write(kademlia.SerialiseGetPacket(&packet))

		header := make([]byte, 4)
		_, err = io.ReadFull(conn, header)
		if err != nil {
			log.Fatal(err)
		}

		size := binary.BigEndian.Uint16(header[0:2])
		messageType := binary.BigEndian.Uint16(header[2:4])
		body := make([]byte, size-4)
		io.ReadFull(conn, body)
		response := append(header, body...)
		switch messageType {
		case kademlia.FAILURE:
			log.Print("No value found")
			conn.Close()

		case kademlia.SUCCESS:

			responsePack, err := kademlia.ParseSucceessPacket(response)
			if err != nil {
				log.Fatal(err)
			}
			log.Printf("Found value %s", responsePack.Value)
			conn.Close()

		default:
			flag.PrintDefaults()
		}

	case "ping":
		cert, err := os.ReadFile(*certPath)
		if err != nil {
			log.Fatal(err)
		}

		ip, port, err := parseIP(*address)
		conn, err := kademlia.ConnectToNode(kademlia.P2PNode{
			NodeID:      [32]byte{},
			IPAddress:   net.ParseIP(ip),
			Port:        port,
			Certificate: cert,
		})
		if err != nil {
			log.Fatal(err)
		}
		conn.Write(kademlia.SerialisePingPacket(&kademlia.P2PPingPacket{
			SourceIP:    net.ParseIP(ip),
			SourcePort:  uint16(port),
			Certificate: cert,
		}))

		reply := make([]byte, 4)
		_, err = io.ReadFull(conn, reply)
		log.Print("Got Reply")

	case "find":
		cert, err := os.ReadFile(*certPath)
		if err != nil {
			log.Fatal(err)
		}

		ip, port, err := parseIP(*address)
		conn, err := kademlia.ConnectToNode(kademlia.P2PNode{
			NodeID:      [32]byte{},
			IPAddress:   net.ParseIP(ip),
			Port:        port,
			Certificate: cert,
		})
		if err != nil {
			log.Fatal(err)
		}
		conn.Write(kademlia.SerialiseFindValuePacket(&kademlia.P2PFindValuePacket{
			SourceIP:    net.ParseIP(ip),
			SourcePort:  uint16(port),
			Certificate: cert,
			Key:         key,
		}))

		responseHeader := make([]byte, 4)
		_, err = io.ReadFull(conn, responseHeader)
		if err != nil {
			log.Print(err)
			return
		}
		size := binary.BigEndian.Uint16(responseHeader[0:2])
		responseType := binary.BigEndian.Uint16(responseHeader[2:4])

		switch responseType {
		case kademlia.FOUNDVALUE:
			responseBody := make([]byte, size-4)
			_, err = io.ReadFull(conn, responseBody)
			if err != nil {
				log.Print(err)
				return
			}
			reply := kademlia.ParseFoundValuePacket(append(responseHeader, responseBody...))
			log.Printf("Found value %s", string(reply.Value))
		default:
			log.Printf("Value not present")
		}
	}

}

func parseIP(line string) (string, uint64, error) {
	var temp []string
	if strings.HasPrefix(line, "[") {
		temp = strings.Split(line[1:], "]:")
	} else {
		temp = strings.Split(line, ":")
	}
	if len(temp) != 2 {
		return "", 0, errors.New("invalid IP format")
	}

	portTemp, err := strconv.Atoi(temp[1])
	if err != nil {
		return "", 0, err
	}
	port := uint64(portTemp)
	return temp[0], port, nil
}