shish.py

import os
import shlex
import signal
import sys
import traceback
import readline

PROMPT = '(%(sh)s)%(user)s@%(host)s %(wd)s%(prch)s '

def prompt():
    return PROMPT % {
        'sh': 'shish',
        'user': 'USER',
        'uid': os.getuid(),
        'host': os.uname()[1],
        'wd': os.getcwd(),
        'prch': '#' if os.geteuid() == 0 else '$',
    }

def sigchld_handler(*args):
    while True:
        try:
            pid, status = os.waitpid(-1, os.WNOHANG)
        except OSError:
            break
        else:
            if pid == 0:
                break
            try:
                JOBS.remove(pid)
            except ValueError:
                sys.stderr.write('<(not a job)>')
            sys.stderr.write('<Process %d exited %d>\n'%(pid, status))

signal.signal(signal.SIGCHLD, sigchld_handler)

JOBS = []

def execute(parts):
    fg = True
    rediri = None
    rediro = None
    oappend = False
    redire = None
    eappend = False
    for part in parts[:]:
        if part.startswith('>>'):
            rediro = part[2:]
            oappend = True
            parts.remove(part)
        elif part.startswith('>'):
            rediro = part[1:]
            parts.remove(part)
        if part.startswith('2>>'):
            redire = part[2:]
            eappend = True
            parts.remove(part)
        elif part.startswith('2>'):
            redire = part[1:]
            parts.remove(part)
        if part.startswith('<'):
            rediri = part[1:]
            parts.remove(part)
        if part == '&':
            fg = False
            parts.remove(part)
    pid = os.fork()
    if not pid:  # Child
        os.dup2(2, 9)
        if rediri:
            fd = os.open(rediri, os.O_RDONLY)
            os.dup2(fd, 0)
            os.close(fd)
        elif not fg:
            os.close(0)
        if rediro:
            fd = os.open(rediro, os.O_WRONLY | (os.O_APPEND if oappend else os.O_TRUNC))
            os.dup2(fd, 1)
            os.close(fd)
        if redire:
            fd = os.open(redire, os.O_WRONLY | (os.O_APPEND if oappend else os.O_TRUNC))
            os.dup2(fd, 2)
            os.close(fd)
        try:
            os.execvp(parts[0], parts)
        except OSError as e:
            os.write(9, ('exec: %s\n'%(e.strerror,)).encode('utf8'))
            os._exit(127)
    else:  # Parent
        if fg:
            pid, status = os.waitpid(pid, 0)
            sys.stderr.write('<%d>'%(status,))
        else:
            sys.stderr.write('<Process %d spawned>\n'%(pid,))
            JOBS.append(pid)

def bt_exit(parts):
    if len(parts) > 1:
        exit(int(parts[1]))
    if JOBS:
        sys.stderr.write('<There are stopped jobs.>\n')
        return
    exit(0)

def bt_jobs(parts):
    for job in JOBS:
        sys.stderr.write('< %d = '%(job,))
        try:
            sys.stderr.write(open('/proc/%d/cmdline'%(job,)).read())
        except Exception as e:
            sys.stderr.write(repr(e))
        sys.stderr.write(' >\n')

def bt_kill(parts):
    pid = int(parts[1])
    sig = 15
    if len(parts) > 2:
        sig = int(parts[2])
    try:
        os.kill(pid, sig)
    except OSError as e:
        sys.stderr.write('kill: %s\n'%(e.strerror,))

def bt_prompt(parts):
    global PROMPT
    PROMPT = parts[1]

readline.parse_and_bind('tab: complete')
while True:
    try:
        line = input(prompt())
        parts = shlex.split(line)
        if not parts:
            continue
        globals().get('bt_'+parts[0], execute)(parts)
    except Exception:
        traceback.print_exc()