package godrop

import (
	"errors"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"syscall"
)

func readFile(name, file string) (id int, err error) {
	id = -1
	if name == "" {
		return
	}
	f, err := os.Open(file)
	if err != nil {
		return
	}
	b, err := ioutil.ReadAll(f)
	if err != nil {
		return
	}
	r := regexp.MustCompile(name + ":.*:(\\d+)")
	m := r.FindAllStringSubmatch(string(b), 1)
	if len(m) == 1 {
		if len(m[0]) == 2 {
			return strconv.Atoi(m[0][1])
		}
	}
	return
}

// UID of user
func UID(user string) (int, error) {
	return readFile(user, "/etc/passwd")
}

// GID of group
func GID(group string) (int, error) {
	return readFile(group, "/etc/group")
}

type Config struct {
	User   string
	Group  string
	Chroot string
}

// Drop will spawn a new process, hand over the listening socket file descriptor and terminate itself after
func Drop(c Config, f func() (net.Listener, error)) error {
	uid, err := UID(c.User)
	if err != nil {
		return err
	}
	gid, err := GID(c.Group)
	if err != nil {
		return err
	}
	switch os.Getuid() {
	case 0:
		ln, err := f()
		if err != nil {
			return err
		}

		l, ok := ln.(*net.TCPListener)
		if !ok {
			return errors.New("godrop: interface conversion failed")
		}
		f, err := l.File()
		if err != nil {
			return err
		}

		cmd := exec.Command(os.Args[0], os.Args[1:]...)
		cmd.ExtraFiles = []*os.File{f}
		cmd.SysProcAttr = &syscall.SysProcAttr{
			Chroot: c.Chroot,
			Credential: &syscall.Credential{
				Uid: uint32(uid),
				Gid: uint32(gid),
			},
			Setsid: true,
		}

		if err := cmd.Start(); err != nil {
			return err
		}

		cmd.Process.Release()
		os.Exit(0)
	case uid:
		return nil
	}
	return errors.New("godrop: droping priviledges failed")
}

// MultiDrop will spawn a new process, hand over the all listening sockets and terminate itself after
func MultiDrop(c Config, f func() ([]net.Listener, error)) error {
	uid, err := UID(c.User)
	if err != nil {
		return err
	}
	gid, err := GID(c.Group)
	if err != nil {
		return err
	}
	switch os.Getuid() {
	case 0:
		cmd := exec.Command(os.Args[0], os.Args[1:]...)
		ln, err := f()
		if err != nil {
			return err
		}
		for _, v := range ln {
			l, ok := v.(*net.TCPListener)
			if !ok {
				return errors.New("godrop: interface conversion failed")
			}
			f, err := l.File()
			if err != nil {
				return err
			}
			cmd.ExtraFiles = append(cmd.ExtraFiles, f)
		}

		cmd.SysProcAttr = &syscall.SysProcAttr{
			Chroot: c.Chroot,
			Credential: &syscall.Credential{
				Uid: uint32(uid),
				Gid: uint32(gid),
			},
			Setsid: true,
		}

		if err := cmd.Start(); err != nil {
			return err
		}

		cmd.Process.Release()
		os.Exit(0)
	case uid:
		return nil
	}
	return errors.New("godrop: droping priviledges failed")
}

// GetListener returns the listener socket of file descriptor 3
func GetListener() (net.Listener, error) {
	return net.FileListener(os.NewFile(3, "[socket]"))
}

// GetListenerFd returns the listener socket of the given file descriptor
func GetListenerFd(fd int) (net.Listener, error) {
	return net.FileListener(os.NewFile(uintptr(fd), "[socket]"))
}