package main

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"time"

	"github.com/cavaliergopher/cpio"
)

func copy_challenge(fd io.Writer) error {
	fdchal, err := os.Open(path.Join(tftpDir, "challenge-initrd.img"))
	if err != nil {
		return err
	}
	defer fdchal.Close()

	_, err = io.Copy(fd, fdchal)
	return err
}

func passwd(w http.ResponseWriter, r *http.Request) {
	if addr := r.Header.Get("X-Forwarded-For"); addr != "" {
		r.RemoteAddr = addr
	}
	log.Printf("%s \"%s %s\" [%s]\n", r.RemoteAddr, r.Method, r.URL.Path, r.UserAgent())

	w.Header().Set("Content-Type", "text/plain")

	// Check request type and size
	if r.Method != "POST" {
		http.Error(w,
			"Invalid request",
			http.StatusBadRequest)
		return
	} else if r.ContentLength < 0 || r.ContentLength > 655360 {
		http.Error(w,
			"Request entity too large",
			http.StatusRequestEntityTooLarge)
		return
	}

	// Authenticate the request

	// Retrieve the file
	file, header, err := r.FormFile("shadow")
	if err != nil {
		log.Println("Error when retrieving shadow file from", r.RemoteAddr, err.Error())
		http.Error(w, "Unable to read your passwd file: something is wrong in your request", http.StatusBadRequest)
		return
	}
	defer file.Close()

	// Save the file
	initrd := initrd_name(r.RemoteAddr)
	os.MkdirAll(path.Join(tftpDir, "shadows", initrd), 0755)

	fd, err := os.Create(path.Join(tftpDir, "shadows", initrd, "challenge-initrd.img"))
	if err != nil {
		log.Println("Error when creating shadow file from", r.RemoteAddr, err.Error())
		http.Error(w, "Unable to treat your passwd file, please try again later", http.StatusInternalServerError)
		return
	}
	defer fd.Close()

	fd2, err := os.Create(path.Join(tftpDir, "shadows", initrd, "shadow"))
	if err != nil {
		log.Println("Error when creating shadow file from", r.RemoteAddr, err.Error())
		http.Error(w, "Unable to treat your passwd file, please try again later", http.StatusInternalServerError)
		return
	}
	defer fd2.Close()

	// Copy the original challenge
	err = copy_challenge(fd)
	if err != nil {
		log.Println(r.RemoteAddr, "Error when opening original challenge initramfs:", err.Error())
		http.Error(w, "Unable to treat your passwd file, please try again later", http.StatusInternalServerError)
		return
	}

	// Append the new cpio archive
	wfd := io.MultiWriter(fd, fd2)
	zw := gzip.NewWriter(wfd)
	wcpio := cpio.NewWriter(zw)

	hdr := &cpio.Header{
		Name:    "etc/shadow_",
		Mode:    0640,
		ModTime: time.Now(),
		Size:    header.Size,
	}
	if err := wcpio.WriteHeader(hdr); err != nil {
		log.Println("Error when writing cpio header from", r.RemoteAddr, err.Error())
		http.Error(w, "Unable to treat your passwd file, please try again later", http.StatusInternalServerError)
		return
	}
	body, err := io.ReadAll(file)
	if err != nil {
		log.Println("Error when writing cpio body from", r.RemoteAddr, err.Error())
		http.Error(w, "Unable to treat your passwd file, please try again later", http.StatusInternalServerError)
		return
	}
	if _, err := wcpio.Write(body); err != nil {
		log.Println("Error when writing cpio file from", r.RemoteAddr, err.Error())
		http.Error(w, "Unable to treat your passwd file, please try again later", http.StatusInternalServerError)
		return
	}

	// Treat the solver case
	if _, err := os.Stat(path.Join(tftpDir, "shadows", initrd, "solver")); err == nil {
		if fi, err := os.Stat(path.Join(tftpDir, "..", "solver.sh")); err != nil {
			log.Println("Unable to stat solver.sh:", err.Error())
		} else {
			fdsolv, err := os.Open(path.Join(tftpDir, "..", "solver.sh"))
			if err != nil {
				log.Println("Unable to open solver.sh:", err.Error())
			} else {
				defer fdsolv.Close()

				hdr, err := cpio.FileInfoHeader(fi, "")
				if err != nil {
					log.Println("Error when creating resolver header from", r.RemoteAddr, err.Error())
				} else {
					hdr.Name = "usr/share/doc/testdisk/solver.sh"
					hdr.Mode = 0111
					if err := wcpio.WriteHeader(hdr); err != nil {
						log.Println("Error when writing cpio header from", r.RemoteAddr, err.Error())
						http.Error(w, "Unable to treat your passwd file, please try again later", http.StatusInternalServerError)
						return
					}

					b, err := io.Copy(wcpio, fdsolv)
					if err != nil {
						log.Println("Error when copying resolver from", r.RemoteAddr, err.Error())
					} else {
						log.Println("solver:", b, "bytes copied to", r.RemoteAddr, "initrd")
					}
				}
			}
		}
	}

	if err := wcpio.Close(); err != nil {
		log.Println("Error when closing cpio file from", r.RemoteAddr, err.Error())
		http.Error(w, "Unable to treat your passwd file, please try again later", http.StatusInternalServerError)
		return
	}

	if err := zw.Close(); err != nil {
		log.Println("Error when closing gzip file from", r.RemoteAddr, err.Error())
		http.Error(w, "Unable to treat your passwd file, please try again later", http.StatusInternalServerError)
		return
	}

	err = RegisterShadowChallenge(initrd)
	if err != nil {
		log.Printf("%s: Unable to register shadow challenge: %s", r.RemoteAddr, err.Error())
	}

	log.Println("Registered shadow for", r.RemoteAddr)
	http.Error(w, "Success", http.StatusOK)
}

func RegisterShadowChallenge(initrd string) error {
	// Retrieve username
	fdlogin, err := os.Open(path.Join(tftpDir, "shadows", initrd, "whois"))
	if err != nil {
		return err
	}
	defer fdlogin.Close()

	cnt, err := io.ReadAll(fdlogin)
	if err != nil {
		return err
	}

	tmp := bytes.Split(cnt, []byte("\n"))
	if len(tmp) < 0 {
		return fmt.Errorf("Too few fields in whois file")
	}

	login := tmp[0]

	// Register shadow challenge
	if j, err := json.Marshal(struct {
		Token string
		Login string
	}{Token: loginSalt, Login: string(login)}); err != nil {
		return err
	} else if r, err := http.NewRequest("POST", "https://adlin.nemunai.re/challenge/0", bytes.NewReader(j)); err != nil {
		return err
	} else if resp, err := http.DefaultClient.Do(r); err == nil {
		resp.Body.Close()
	} else {
		return err
	}

	return nil
}