midutils/midsu.py

import argparse
import configparser
import os
import sys
import stat
import getpass
import pwd
import grp
import subprocess
import re
import fcntl
import shutil

CONFIG_FILE = '/etc/midsu'
DEFAULT_SHELL = '/usr/bin/tmux', '/bin/bash', '/bin/sh'
SHELL_OVERRIDE_FILE = '.midsu'
SHADOW_HOMEDIR = '.shadow-home'
SKELETON = '/etc/skel'
TMP_DIR = '/tmp'
TMP_PREFIX = 'midsu-'

# Expected executable paths
PATH = {
    'bwrap': '/usr/bin/bwrap',
    'python3': '/usr/bin/python3',
    'runuser': '/sbin/runuser',
    'sudo': '/usr/bin/sudo',
    'su': '/bin/su',
    'useradd': '/sbin/useradd'
}


def read_config(config_file: str) -> configparser.ConfigParser:
    """Read and validate global midsu config
    """
    
    config = configparser.ConfigParser()
    
    # Load config file
    if not os.path.isfile(config_file):
        raise FileNotFoundError(f"Global config file ({config_file}) not found")
    config.read(config_file)
    
    # Validate primary-shadow mapping
    mapping = dict(config['shadows'])
    
    primary_users = mapping.keys()
    shadow_users = mapping.values()
    if len(set(shadow_users)) != len(shadow_users):
        raise ValueError("Invalid mapping: duplicate shadow users found")
    
    intersection = set(primary_users) & set(shadow_users)
    if intersection:
        raise ValueError("Invalid mapping: some users are both primary and shadow")
    
    return config

def get_current_user() -> str:
    """Return invoking user if UID is 0
    """
    
    if os.getuid() == 0:
        #user = os.getlogin()
        user = os.environ.get('SUDO_USER')
    else:
        user = getpass.getuser()
    
    return user

def get_current_mapping(user: str, mapping: dict) -> tuple:
    """Return primary user and shadow user
    """
    
    if user in mapping.keys():
        primary_user = user
        shadow_user = mapping.get(primary_user)
    elif user in mapping.values():
        shadow_user = user
        for key, value in mapping.items():
            if value == shadow_user:
                primary_user = key
                break
    # If user not in mapping
    else:
        primary_user = shadow_user = None
    
    return primary_user, shadow_user

def user_exists(user: str) -> bool:
    """Return True if user exist or False otherwise
    """
    
    try:
        pwd.getpwnam(user)
        return True
    except KeyError:
        return False

STRIP_ARGV_PATTERN = re.compile(r'^-+')
def strip_argv(arg: str) -> str:
    """Remove unsafe characters for arg from string
    """
    
    return STRIP_ARGV_PATTERN.sub('', arg)

def create_shadow_user(primary_user: str):
    """Create new shadow user for given primary user
    """
    
    primary_user_info = pwd.getpwnam(primary_user)
    primary_homedir = primary_user_info.pw_dir
    
    # Determine new user details
    shadow_user = primary_user + '-shadow'
    shadow_homedir = os.path.join(primary_homedir, SHADOW_HOMEDIR)
    user_group = grp.getgrgid(primary_user_info.pw_gid).gr_name
    
    # If username in use, look for different one
    if user_exists(shadow_user):
        i = 1
        while user_exists(shadow_user + str(i)):
            i += 1
        shadow_user = shadow_user + str(i)
    
    # Validate bad names
    if shadow_user.startswith('-'):
        raise ValueError(f"Invalid username '{shadow_user}'")
    
    # Exit if shadow homedir exists
    if os.path.exists(shadow_homedir):
        raise FileExistsError(f"Shadow user homedir ({shadow_homedir}) already exist")
    
    # Create the user using 'useradd' command
    cmd_args = (
        '--system',
        '--home-dir', strip_argv(shadow_homedir),
        '--no-create-home',
        '--gid',      strip_argv(user_group),
                      strip_argv(shadow_user)
    )
    try:
        subprocess.run((PATH['useradd'], *cmd_args), check=True)
    except subprocess.CalledProcessError as errmsg:
        print("Failed to create shadow user:", errmsg)
        return
    
    # Lock down primary homedir
    shadow_user_info = pwd.getpwnam(shadow_user)
    uid, gid = shadow_user_info.pw_uid, shadow_user_info.pw_gid
    os.chown(primary_homedir, uid, gid)
    primary_mode = os.stat(primary_homedir).st_mode | stat.S_ISVTX | stat.S_IRGRP | stat.S_IXGRP
    os.chmod(primary_homedir, primary_mode)
    
    # Create homedir
    run_as(shadow_user)
    shutil.copytree(SKELETON, shadow_homedir, dirs_exist_ok=False)
    shadow_mode = stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR
    os.chmod(shadow_homedir, shadow_mode)
    run_as()
    
    return shadow_user

def get_login_program(user: str, failsafe: bool = False) -> str:
    """Determine midsu login program for user
    """
    
    def default():
        for shell in DEFAULT_SHELL:
            if os.access(shell, os.X_OK, effective_ids=True):
                return shell
    
    homedir = pwd.getpwnam(user).pw_dir
    override_file = os.path.join(homedir, SHELL_OVERRIDE_FILE)
    
    # Return default program in failsafe mode
    if failsafe:
        return default()
    
    # Return custom program if exists
    if os.path.isfile(override_file):
        return override_file
    
    # Fallback to default program
    return default()

def switch_user(user: str, failsafe: bool = False):
    """Exec midsu login program as given shadow user
    """
    
    # Validate username
    if user.startswith('-'):
        raise ValueError(f"Invalid username '{user}'")
    
    if not user_exists(user):
        raise ValueError(f"User {user} does not exist")
    
    # Get the login program
    login_program = get_login_program(user, failsafe=failsafe)
    
    os.execv(PATH['su'],
        ('su', strip_argv(user), '-s', strip_argv(login_program))
    )

def restart_with_sudo(__file__: str):
    """Invoke itself with sudo
    """
    
    if os.getuid() != 0:
        path_to_self = os.path.abspath(__file__)
        exec_args = sys.argv
        exec_args[0] = path_to_self
        os.execv(PATH['sudo'],
            ('sudo', '--', PATH['python3'], '-E', *exec_args)
        )

def run_as(username: str = 'root', lock: bool = False):
    """Changes the process' uid. Optionally locks it down, for it cannot be reversed
    """
    
    # Exit if switching is locked
    if 0 not in os.getresuid():
        raise PermissionError(f"Switching users is locked")
    
    # Determine target uid and gid
    if username == 'root':
        t_uid = t_gid = 0
    else:
        user_info = pwd.getpwnam(username)
        t_uid = user_info.pw_uid
        t_gid = user_info.pw_gid
    
    os.seteuid(0)
    
    # Clear groups
    os.setgroups([])
    os.setgid(t_gid)
    
    # Set new uid
    if lock:
        os.setuid(t_uid)
    else:
        os.seteuid(t_uid)

def main():
    # Parse command-line arguments
    parser = argparse.ArgumentParser(
        description="Switch to your shadow user"
    )
    parser.add_argument(
        '--failsafe',
        action='store_true',
        help=f"Run default program {DEFAULT_SHELL}"
    )
    args = parser.parse_args()
    
    if os.getuid() == 0:
        # Lock config
        fd = open(CONFIG_FILE, 'a')
        fcntl.lockf(fd, fcntl.LOCK_EX)
    
    # Read config
    config = read_config(CONFIG_FILE)
    
    # Get current, primary and shadow username
    current_user = get_current_user()
    primary_user, shadow_user = get_current_mapping(current_user, dict(config['shadows']))
    
    # Exit if already logged in as shadow user
    if current_user == shadow_user:
        print("Already logged in as shadow user")
        exit(1)
    
    # Exit if user is unmapped and autocreating user not permitted
    user_auto_create = config.getboolean('main', 'user_auto_create', fallback=True)
    if not shadow_user and not user_auto_create:
        print("You don't have shadow user to switch to")
        exit(1)
    
    restart_with_sudo(__file__)
    
    # If shadow user does not exist, create it
    if not shadow_user:
        print("Creating your shadow user...")
        shadow_user = create_shadow_user(current_user)
        # Exit on error
        if not shadow_user:
            return
        primary_user = current_user
        
        # Update mapping
        config.set('shadows', primary_user, shadow_user)
        with open(CONFIG_FILE, 'w') as fd:
            config.write(fd)
        
        print("Shadow user created successfully!")
    
    # Log in as shadow user
    switch_user(shadow_user, failsafe=args.failsafe)


if __name__ == '__main__':
    main()