ch3.cpp

#include <iostream>
#include <cstdio>
#include <cstddef>
#include <csignal>
#include <cstdlib>
#include <unistd.h>
#include <sys/wait.h>

#include <string>
#include <vector>
#include <list>
#include <algorithm>
#include <thread>

#ifndef _COLORS_ // From http://stackoverflow.com/questions/2616906/how-do-i-output-coloured-text-to-a-linux-terminal
#define _COLORS_

/* FOREGROUND */
#define RST  "\x1B[0m"
#define KRED  "\x1B[31m"
#define KGRN  "\x1B[32m"
#define KYEL  "\x1B[33m"
#define KBLU  "\x1B[34m"
#define KMAG  "\x1B[35m"
#define KCYN  "\x1B[36m"
#define KWHT  "\x1B[37m"

#define FRED(x) KRED x RST
#define FGRN(x) KGRN x RST
#define FYEL(x) KYEL x RST
#define FBLU(x) KBLU x RST
#define FMAG(x) KMAG x RST
#define FCYN(x) KCYN x RST
#define FWHT(x) KWHT x RST

#define BOLD(x) "\x1B[1m" x RST
#define UNDL(x) "\x1B[4m" x RST

#endif  /* _COLORS_ */

std::vector<pid_t> pids;

static const std::string WELCOME_STR = "Welcome to htShell!";

struct ParsedCmd
{
    std::vector<std::string> args;
    bool isBackground;
};

struct ParsedInput
{
    std::vector<ParsedCmd> pipedCmd;
};

ParsedCmd parseCmd(std::string cmd)
{
    ParsedCmd pc;
    pc.isBackground = false;
    std::string cur;
    for (std::size_t i=0; i<cmd.size(); ++i)
    {
        if (cmd[i] == ' ' || cmd[i] == '&')
        {
            if (!cur.empty())
            {
                pc.args.push_back(cur);
                cur.clear();
            }
            if (cmd[i] == '&')
                pc.isBackground = true;
        }
        else
            cur += cmd[i];
    }
    if (!cur.empty())
        pc.args.push_back(cur);
    return pc;

}

ParsedInput parseInput(std::string cmd)
{
    std::string cur;
    ParsedInput pi;
    for (std::size_t i=0; i<cmd.size(); ++i)
    {
        if (cmd[i] == '|')
        {
            if (std::count(cur.cbegin(), cur.cend(), ' ') != cur.length())
                pi.pipedCmd.push_back(parseCmd(cur));
            cur = "";
        }
        else
            cur += cmd[i];
    }
    if (std::count(cur.cbegin(), cur.cend(), ' ') != cur.length())
        pi.pipedCmd.push_back(parseCmd(cur));
    return pi;
}

void setup()
{
    std::cout << BOLD(<<WELCOME_STR<<) << std::endl;
}

std::string readInput()
{
    std::string ans;
    getline(std::cin, ans);
    return ans;
}

void handle_ch_SIGINT(int signal)
{
    exit(EXIT_FAILURE);
}


struct Task
{
    pid_t tpid;
    std::string input;
    bool isBackground;
    std::thread thr;
};
std::list<Task> tasks;


void task_watch(pid_t tpid, std::string input, bool isBackground)
{
    Task t { tpid, input, isBackground};
    auto it = tasks.insert(tasks.cbegin(), std::move(t));

    it->thr = std::thread([tpid, isBackground, it]() {
                    waitpid(tpid, NULL, 0);
                    if (isBackground)
                    tasks.erase(it);
                    });
    if (isBackground)
        it->thr.detach();
    if (!isBackground)
    {
        it->thr.join();
        tasks.erase(it);
    }
}
std::vector<std::string> history;

void handleInput(ParsedInput pi, std::string raw)
{
    // handle null input
    if (pi.pipedCmd.empty())
        return;

    if (pi.pipedCmd.size() == 1)
    {
        ParsedCmd pc = pi.pipedCmd[0];
        if (pc.args[0] == "r")
        {
            switch(pc.args.size())
            {
                case 1:
                    {
                        if (!history.empty())
                        {
                            pi = parseInput(*history.rbegin());
                            raw = *history.rbegin();
                            break;
                        }
                        else
                        {
                            std::cout << "No history!" << std::endl;
                            return;
                        }
                    }
                case 2:
                    {
                        std::string to_find = pc.args[1];
                        auto it = std::find_if(history.rbegin(), history.rend(), [to_find](const std::string& s) {
                                    return s.find(to_find) == 0;
                                    });
                        if (it == history.rend())
                        {
                            std::cout << "Not found history that begins with " << to_find << std::endl;
                            return;
                        }
                        pi = parseInput(*it);
                        raw = *it;
                        break;
                    }
            }
        }
        if (pc.args[0] == "cd" && pc.args.size() == 2)
        {
            if (chdir(pc.args[1].c_str()) < 0)
                perror("cd");
            return;
        }
        if (pc.args[0] == "jobs" && pc.args.size() == 1)
        {
            int cnt = 0;
            std::for_each(tasks.begin(), tasks.end(), [&cnt](const Task& ta)
                        {
                        std::cout << cnt++ << "\t" << "[" << ta.tpid  << "]" << ta.input << std::endl;
                        });
            return;
        }
        if (pc.args[0] == "fg" && pc.args.size() == 2)
        {
            int gpid = std::atoi(pc.args[1].c_str());

            pid_t stdinpgrp = tcgetpgrp(STDIN_FILENO);
            signal(SIGTTOU, SIG_IGN);
            if (tcsetpgrp(STDIN_FILENO, gpid))
            {
                perror("tcsetpgrp: stdin");
            }
            if (kill(-gpid, SIGCONT))
                perror("kill");
            waitpid(gpid, NULL, 0);
            if (tcsetpgrp(STDIN_FILENO, stdinpgrp))
            {
                perror("tcsetpgrp: restore");
            }
            signal(SIGTTOU, SIG_DFL);
            return;
        }
    }

    history.push_back(raw);

    pid_t stdinpgrp = tcgetpgrp(STDIN_FILENO);

    pid_t gpid;
    if (gpid = fork()) // top-level ch3 program
    {
        setpgid(gpid, gpid);
        signal(SIGTTOU, SIG_IGN);
        if (tcsetpgrp(STDIN_FILENO, gpid))
        {
            perror("tcsetpgrp: stdin");
        }
        bool isBackground = false;
        std::for_each(pi.pipedCmd.begin(), pi.pipedCmd.end(), [&](const ParsedCmd& pc){
                    isBackground |= pc.isBackground;
                    });
        task_watch(gpid, pi.pipedCmd[0].args[0], isBackground);
        if (tcsetpgrp(STDIN_FILENO, stdinpgrp))
        {
            perror("tcsetpgrp: restore");
        }
        signal(SIGTTOU, SIG_DFL);
    } else
    {   
        pids.clear();
        // Process in charge of whole task
        setpgid(0, 0);
        pid_t pgid = getpgid(0);
        bool isBackground = false;
        std::vector<std::pair<int, int> > pipes;
        for (std::size_t i=0; i<pi.pipedCmd.size() -1; ++i)
        {
            int flides[2];
            pipe(flides);
            pipes.push_back(std::make_pair(flides[1], flides[0]));
        }

        for (std::size_t i=0; i<pi.pipedCmd.size(); ++i)
        {
            ParsedCmd cmd = pi.pipedCmd[i];
            isBackground |= cmd.isBackground;
            char* args_arr[1024];
            for (std::size_t j=0; j<cmd.args.size(); ++j)
                args_arr[j] = const_cast<char*>(cmd.args[j].c_str());
            args_arr[cmd.args.size()] = NULL;
            if (cmd.args.empty())
                continue;

            pid_t pid2;
            if (pid2 = fork())
            {
                pids.push_back(pid2);
            } else
            {   // Process of each command
                signal(SIGTTOU, SIG_DFL);
                setpgid(0, pgid);
                if (i>0)
                    dup2(pipes[i-1].second, STDIN_FILENO);
                if (i<pi.pipedCmd.size() -1)
                {
                    dup2(pipes[i].first, STDOUT_FILENO);
                }
                // Close unused pipes
                for (std::size_t j=0; j<pipes.size(); ++j)
                {
                    if (i-1 != j)
                        close(pipes[j].second);
                    if (i!=j)
                        close(pipes[j].first);
                }
                if (execvp(args_arr[0], args_arr) < 0)
                    perror("Execvp");
            }
        }
        int cnt = pids.size();
        std::vector<bool> is_closed(pipes.size());
        while (cnt > 0)
        {
            int statloc;
            pid_t chpid = waitpid(0, &statloc, 0);
            int idx = std::find(pids.begin(), pids.end(), chpid) - pids.begin();
            if (idx>0 && !is_closed[idx-1])
            {
                is_closed[idx-1] = true;
                if (close(pipes[idx-1].first))
                    perror("close pipe");
                if (close(pipes[idx-1].second))
                    perror("close pipe");

            }
            if (idx < pids.size() - 1 && !is_closed[idx])
            {
                is_closed[idx] = true;
                if ( close(pipes[idx].first) )
                    perror("close pipe");
                if (close(pipes[idx].second))
                    perror("close pipe");
            }
            --cnt;
            if (WIFEXITED(statloc))
            {
                int ret = WEXITSTATUS(statloc);
                if (ret == 0)
                    std::cout << FBLU(<<"[pid="<<chpid<<"] exited with code " << WEXITSTATUS(statloc) <<)<< std::endl;
                else
                    std::cout << FRED(<<"[pid="<<chpid<<"] exited with code " << WEXITSTATUS(statloc) <<)<< std::endl;
            }
            if (WIFSIGNALED(statloc))
                std::cout << FYEL(<<"[pid="<<chpid<<"] caught a signal " << WTERMSIG(statloc) <<) << std::endl;
        }
        exit(0);
    }
}

void handle_default(int)
{
    int cnt = history.size() < 10 ? history.size() : 10;

    std::cout << std::endl;
    std::cout << "History:" << std::endl;
    for (std::size_t i=0; i<cnt; ++i)
    {
        std::cout << i << ". " << history[history.size() - 1 - i] << std::endl;
    }
}

int main()
{
    const std::size_t MAX_CWDBUF = 4096;
    char cwdbuf[MAX_CWDBUF];

    signal(SIGINT, handle_default);

    setup();
    while (!std::cin.eof())
    {
        char *pCwd = getcwd(cwdbuf, MAX_CWDBUF);
        std::fflush(stdout);
        if (!pCwd)
            std::perror("getCwd");
        else
            std::cout << BOLD( << pCwd <<) <<  " >> ";
        std::string input = readInput();
        ParsedInput pc = parseInput(input);
        handleInput(pc, input);
    }
}