// Copyright (C) 2017 Marius Schellenberger

// Package godrop provides a simple privileges dropping library
package godrop

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

func errf(err error) error {
	return fmt.Errorf("godrop: %s", err)
}

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")
}

// Config represents the drop config
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 {
	return MultiDrop(c, func() ([]net.Listener, error) {
		l, err := f()
		return []net.Listener{l}, err
	})
}

// 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 errf(err)
	}
	gid, err := GID(c.Group)
	if err != nil {
		return errf(err)
	}
	switch os.Getuid() {
	case 0:
		cmd := exec.Command(os.Args[0], os.Args[1:]...)
		ln, err := f()
		if err != nil {
			return errf(err)
		}
		for i, v := range ln {
			var f *os.File
			switch l := v.(type) {
			case *net.TCPListener:
				f, err = l.File()
				l.Close()
			case *net.UnixListener:
				f, err = l.File()
				l.Close()
			default:
				return errf(fmt.Errorf("index: %d listener is not type of either *net.TCPListener or *net.UnixListener", i))
			}
			if err != nil {
				return errf(fmt.Errorf("index: %d %s", i, 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 errf(err)
		}

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

// GetListener returns the listener socket of file descriptor 3
func GetListener() (net.Listener, error) {
	return GetListenerFd(3)
}

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