package main

import (
	"fmt"
	"os"
	"os/user"
	"strconv"
	"syscall"
)

// Drop root privileges.
func dropPrivileges(userToSwitchTo string) {

	// Check if we have root privileges.
	var rootUser = false
	var rootPrimaryGroup = false
	var rootSupplementaryGroup = false

	// Check UID and EUID.
	if syscall.Getuid() == 0 || syscall.Geteuid() == 0 {
		rootUser = true
	}
	// Check GID and EGID.
	if syscall.Getgid() == 0 || syscall.Getegid() == 0 {
		rootPrimaryGroup = true
	}
	// Check supplementary groups.
	groups, err := syscall.Getgroups()
	if err != nil {
		fatalError(err)
	}
	for _, groupID := range groups {
		if groupID == 0 {
			rootSupplementaryGroup = true
			break
		}
	}

	if userToSwitchTo == "" {

		if rootUser {
			fmt.Println("When running as root, use the -u option to switch to an unprivileged user.")
			os.Exit(1)
		} else if rootPrimaryGroup || rootSupplementaryGroup {
			fatalError("The user running the program is in the root group;\n" +
				"use the -u option to switch to an unprivileged user.")
		}

	} else { // userToSwitchTo != ""

		// Get user and group IDs for the user we want to switch to.
		userInfo, err := user.Lookup(userToSwitchTo)
		if err != nil {
			fatalError(err)
		}
		// Convert group id and user id from string to int.
		gid, err := strconv.Atoi(userInfo.Gid)
		if err != nil {
			fatalError(err)
		}
		uid, err := strconv.Atoi(userInfo.Uid)
		if err != nil {
			fatalError(err)
		}
		// If the user we want to switch to has root privileges, stop execution.
		if uid == 0 || gid == 0 {
			fatalError("Running as root is not allowed.")
		}
		// Unset supplementary group IDs.
		err = syscall.Setgroups([]int{})
		if err != nil {
			fmt.Fprintln(os.Stderr,
				"Failed to unset supplementary group IDs: "+err.Error())
			if rootSupplementaryGroup {
				fatalError("Failed to drop root privileges. Exiting...")
			}
		}
		// Set group ID (real and effective).
		err = syscall.Setgid(gid)
		if err != nil {
			fmt.Fprintln(os.Stderr, "Failed to set group ID: "+err.Error())
			if rootPrimaryGroup {
				fatalError("Failed to drop root privileges. Exiting...")
			}
		}
		// Set user ID (real and effective).
		err = syscall.Setuid(uid)
		if err != nil {
			fmt.Fprintln(os.Stderr, "Failed to set user ID: "+err.Error())
			if rootUser {
				fatalError("Failed to drop root privileges. Exiting...")
			}
		}

	}

}