package main

import (
	"archive/tar"
	"compress/bzip2"
	"compress/gzip"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"git.wntrmute.dev/kyle/goutils/die"
	"git.wntrmute.dev/kyle/goutils/fileutil"
)

var (
	preserveOwners bool
	preserveMode   bool
	verbose        bool
)

func setupFile(hdr *tar.Header, file *os.File) error {
	if preserveMode {
		if verbose {
			fmt.Printf("\tchmod %0#o\n", hdr.Mode)
		}
		err := file.Chmod(os.FileMode(hdr.Mode & 0xFFFFFFFF)) // #nosec G115
		if err != nil {
			return err
		}
	}

	if preserveOwners {
		fmt.Printf("\tchown %d:%d\n", hdr.Uid, hdr.Gid)
		err := file.Chown(hdr.Uid, hdr.Gid)
		if err != nil {
			return err
		}
	}

	return nil
}

func linkTarget(target, top string) string {
	if filepath.IsAbs(target) {
		return target
	}

	return filepath.Clean(filepath.Join(top, target))
}

// safeJoin joins base and elem and ensures the resulting path does not escape base.
func safeJoin(base, elem string) (string, error) {
	cleanBase := filepath.Clean(base)
	joined := filepath.Clean(filepath.Join(cleanBase, elem))

	absBase, err := filepath.Abs(cleanBase)
	if err != nil {
		return "", err
	}
	absJoined, err := filepath.Abs(joined)
	if err != nil {
		return "", err
	}
	rel, err := filepath.Rel(absBase, absJoined)
	if err != nil {
		return "", err
	}
	if rel == ".." || strings.HasPrefix(rel, ".."+string(os.PathSeparator)) {
		return "", fmt.Errorf("path traversal detected: %s escapes %s", elem, base)
	}
	return joined, nil
}

func handleTypeReg(tfr *tar.Reader, hdr *tar.Header, filePath string) error {
	file, err := os.Create(filePath)
	if err != nil {
		return err
	}
	defer file.Close()

	if _, err = io.Copy(file, tfr); err != nil {
		return err
	}
	return setupFile(hdr, file)
}

func handleTypeLink(hdr *tar.Header, top, filePath string) error {
	file, err := os.Create(filePath)
	if err != nil {
		return err
	}
	defer file.Close()

	srcPath, err := safeJoin(top, hdr.Linkname)
	if err != nil {
		return err
	}
	source, err := os.Open(srcPath)
	if err != nil {
		return err
	}
	defer source.Close()

	if _, err = io.Copy(file, source); err != nil {
		return err
	}
	return setupFile(hdr, file)
}

func handleTypeSymlink(hdr *tar.Header, top, filePath string) error {
	if !fileutil.ValidateSymlink(hdr.Linkname, top) {
		return fmt.Errorf("symlink %s is outside the top-level %s", hdr.Linkname, top)
	}
	path := linkTarget(hdr.Linkname, top)
	if ok, err := filepath.Match(top+"/*", filepath.Clean(path)); !ok {
		return fmt.Errorf("symlink %s isn't in %s", hdr.Linkname, top)
	} else if err != nil {
		return err
	}
	return os.Symlink(linkTarget(hdr.Linkname, top), filePath)
}

func handleTypeDir(hdr *tar.Header, filePath string) error {
	return os.MkdirAll(filePath, os.FileMode(hdr.Mode&0xFFFFFFFF)) // #nosec G115
}

func processFile(tfr *tar.Reader, hdr *tar.Header, top string) error {
	if verbose {
		fmt.Println(hdr.Name)
	}

	filePath, err := safeJoin(top, hdr.Name)
	if err != nil {
		return err
	}

	switch hdr.Typeflag {
	case tar.TypeReg:
		return handleTypeReg(tfr, hdr, filePath)
	case tar.TypeLink:
		return handleTypeLink(hdr, top, filePath)
	case tar.TypeSymlink:
		return handleTypeSymlink(hdr, top, filePath)
	case tar.TypeDir:
		return handleTypeDir(hdr, filePath)
	}
	return nil
}

var compression = map[string]bool{
	"gzip":  false,
	"bzip2": false,
}

type bzipCloser struct {
	r io.Reader
}

func (brc *bzipCloser) Read(p []byte) (int, error) {
	return brc.r.Read(p)
}

func (brc *bzipCloser) Close() error {
	return nil
}

func newBzipCloser(r io.ReadCloser) (io.ReadCloser, error) {
	br := bzip2.NewReader(r)
	return &bzipCloser{r: br}, nil
}

var compressFuncs = map[string]func(io.ReadCloser) (io.ReadCloser, error){
	"gzip":  func(r io.ReadCloser) (io.ReadCloser, error) { return gzip.NewReader(r) },
	"bzip2": newBzipCloser,
}

func verifyCompression() bool {
	var compressed bool
	for _, v := range compression {
		if compressed && v {
			return false
		}
		compressed = compressed || v
	}
	return true
}

func getReader(r io.ReadCloser) (io.ReadCloser, error) {
	for c, v := range compression {
		if v {
			return compressFuncs[c](r)
		}
	}

	return r, nil
}

func openArchive(path string) (io.ReadCloser, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}

	r, err := getReader(file)
	if err != nil {
		return nil, err
	}

	return r, nil
}

var compressFlags struct {
	z bool
	j bool
}

func parseCompressFlags() error {
	if compressFlags.z {
		compression["gzip"] = true
	}

	if compressFlags.j {
		compression["bzip2"] = true
	}

	if !verifyCompression() {
		return errors.New("multiple compression formats specified")
	}

	return nil
}

func usage(w io.Writer) {
	fmt.Fprintf(w, `ChromeOS untar

This is a tool that is intended to support untarring on SquashFS file
systems. In particular, every time it encounters a hard link, it
will just create a copy of the file.

Usage: cruntar [-jmvpz] archive [dest]

Flags:
	-a	Shortcut for -m -p: preserve owners and file mode.
	-j	The archive is compressed with bzip2.
	-m	Preserve file modes.
	-p	Preserve ownership.
	-v	Print the name of each file as it is being processed.
	-z	The archive is compressed with gzip.
`)
}

func init() {
	flag.Usage = func() { usage(os.Stderr) }
}

func main() {
	var archive, help bool
	flag.BoolVar(&archive, "a", false, "Shortcut for -m -p: preserve owners and file mode.")
	flag.BoolVar(&help, "h", false, "print a help message")
	flag.BoolVar(&compressFlags.j, "j", false, "bzip2 compression")
	flag.BoolVar(&preserveMode, "m", false, "preserve file modes")
	flag.BoolVar(&preserveOwners, "p", false, "preserve ownership")
	flag.BoolVar(&verbose, "v", false, "verbose mode")
	flag.BoolVar(&compressFlags.z, "z", false, "gzip compression")
	flag.Parse()

	if help {
		usage(os.Stdout)
		os.Exit(0)
	}

	if archive {
		preserveMode = true
		preserveOwners = true
	}

	err := parseCompressFlags()
	die.If(err)

	if flag.NArg() == 0 {
		return
	}

	top := "./"
	if flag.NArg() > 1 {
		top = flag.Arg(1)
	}

	r, err := openArchive(flag.Arg(0))
	die.If(err)

	tfr := tar.NewReader(r)
	var hdr *tar.Header
	for {
		hdr, err = tfr.Next()
		if errors.Is(err, io.EOF) {
			break
		}
		die.If(err)

		err = processFile(tfr, hdr, top)
		die.If(err)
	}

	r.Close()
}