inherit parent's parameters in new processes

Command line parameters, that is.

4 files changed, 375 insertions, 2 deletions

diff --git a/src/cmd_line.hpp b/src/cmd_line.hpp
new file mode 100644
index 0000000..3a65ec7
--- /dev/null
+++ b/src/cmd_line.hpp
@@ -0,0 +1,194 @@
+#pragma once
+
+#include "error.hpp"
+#include "string.hpp"
+
+#include <Windows.h>
+#include <shellapi.h>
+
+#include <cstddef>
+
+#include <memory>
+#include <stdexcept>
+#include <sstream>
+#include <string>
+#include <utility>
+#include <vector>
+
+class CommandLine
+{
+public:
+    static CommandLine query()
+    {
+        return build_from_string(GetCommandLine());
+    }
+
+    static CommandLine build_from_main(int argc, wchar_t* argv[])
+    {
+        if (argc < 1)
+            throw std::range_error(__FUNCTION__ ": invalid argc value");
+
+        std::wstring argv0{argv[0]};
+        --argc;
+        ++argv;
+
+        std::vector<std::wstring> args;
+        args.reserve(argc);
+
+        for (int i = 0; i < argc; ++i)
+            args.emplace_back(argv[i]);
+
+        return {std::move(argv0), std::move(args)};
+    }
+
+    CommandLine() = default;
+
+    bool has_argv0() const
+    {
+        return !argv0.empty();
+    }
+
+    std::wstring get_argv0() const
+    {
+        return argv0;
+    }
+
+    std::wstring escape_argv0() const
+    {
+        return escape(get_argv0());
+    }
+
+    bool has_args() const
+    {
+        return !get_args().empty();
+    }
+
+    const std::vector<std::wstring>& get_args() const
+    {
+        return args;
+    }
+
+    std::vector<std::wstring> escape_args() const
+    {
+        std::vector<std::wstring> safe;
+        safe.reserve(args.size());
+        for (const auto& arg : args)
+            safe.emplace_back(escape(arg));
+        return safe;
+    }
+
+    static constexpr auto sep = L' ';
+
+    std::wstring join_args() const
+    {
+        return string::join(sep, escape_args());
+    }
+
+    std::wstring join() const
+    {
+        if (!has_argv0())
+            throw std::logic_error(__FUNCTION__ ": doesn't have executable path");
+        std::wostringstream oss;
+        oss << escape_argv0();
+        if (has_args())
+            oss << sep << string::join(sep, escape_args());
+        return oss.str();
+    }
+
+private:
+    static CommandLine build_from_string(std::wstring src)
+    {
+        string::trim(src);
+        if (src.empty())
+            return {};
+
+        int argc = 0;
+        std::unique_ptr<wchar_t*, LocalDelete> argv{CommandLineToArgvW(src.c_str(), &argc)};
+
+        if (argv.get() == NULL)
+            error::raise("CommandLineToArgvW");
+
+        if (argc == 0)
+            return {};
+
+        std::wstring argv0{argv.get()[0]};
+
+        std::vector<std::wstring> args;
+        args.reserve(argc - 1);
+
+        for (int i = 1; i < argc; ++i)
+            args.emplace_back(argv.get()[i]);
+
+        return {std::move(argv0), std::move(args)};
+    }
+
+    inline std::wstring escape_for_cmd(const std::wstring& arg)
+    {
+        static constexpr auto escape_symbol = L'^';
+        static constexpr auto dangerous_symbols = L"!\"%&()<>^|";
+
+        auto safe = escape(arg);
+        string::prefix_with(safe, dangerous_symbols, escape_symbol);
+        return safe;
+    }
+
+    static std::wstring escape(const std::wstring& arg)
+    {
+        std::wstring safe;
+        safe.reserve(arg.length() + 2);
+
+        safe.push_back(L'"');
+
+        for (auto it = arg.cbegin(); it != arg.cend(); ++it)
+        {
+            std::size_t numof_backslashes = 0;
+
+            for (; it != arg.cend() && *it == L'\\'; ++it)
+                ++numof_backslashes;
+
+            if (it == arg.cend())
+            {
+                safe.reserve(safe.capacity() + numof_backslashes);
+                safe.append(2 * numof_backslashes, L'\\');
+                break;
+            }
+
+            switch (*it)
+            {
+                case L'"':
+                    safe.reserve(safe.capacity() + numof_backslashes + 1);
+                    safe.append(2 * numof_backslashes + 1, L'\\');
+                    break;
+
+                default:
+                    safe.append(numof_backslashes, L'\\');
+                    break;
+            }
+
+            safe.push_back(*it);
+        }
+
+        safe.push_back(L'"');
+        return safe;
+    }
+
+    struct LocalDelete
+    {
+        void operator()(wchar_t* argv[]) const
+        {
+            LocalFree(argv);
+        }
+    };
+
+    CommandLine(std::vector<std::wstring>&& args)
+        : args{std::move(args)}
+    { }
+
+    CommandLine(std::wstring&& argv0, std::vector<std::wstring>&& args = {})
+        : argv0{std::move(argv0)}
+        , args{std::move(args)}
+    { }
+
+    const std::wstring argv0;
+    const std::vector<std::wstring> args;
+};
diff --git a/src/main.cpp b/src/main.cpp
index 3ae2bcc..9d897b7 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -169,7 +169,7 @@ void on_button_elevate_click(HWND wnd)
 
     try
     {
-        process::runas(process::get_executable_path(), wnd);
+        process::runas_self(wnd);
     }
     catch (const Error& e)
     {
diff --git a/src/process.hpp b/src/process.hpp
index 8279c3c..6fbaf8a 100644
--- a/src/process.hpp
+++ b/src/process.hpp
@@ -5,6 +5,7 @@
 
 #pragma once
 
+#include "cmd_line.hpp"
 #include "error.hpp"
 
 #include <Windows.h>
@@ -29,17 +30,41 @@ namespace process
         return buf.data();
     }
 
-    void runas(const std::wstring& exe_path, HWND hwnd = NULL, int nShow = SW_NORMAL)
+    std::wstring get_command_line()
     {
+        return GetCommandLine();
+    }
+
+    void runas(
+        const CommandLine& cmd_line,
+        HWND hwnd = NULL,
+        int nShow = SW_NORMAL)
+    {
+        static constexpr auto sep = L' ';
+
+        const auto exe_path = cmd_line.has_argv0()
+            ? cmd_line.get_argv0()
+            : get_executable_path();
+
         SHELLEXECUTEINFOW info;
         ZeroMemory(&info, sizeof(info));
         info.cbSize = sizeof(info);
         info.lpVerb = L"runas";
         info.lpFile = exe_path.c_str();
+        const auto args = cmd_line.join_args();
+        if (!args.empty())
+            info.lpParameters = args.c_str();
         info.hwnd = hwnd;
         info.nShow = nShow;
 
         if (!ShellExecuteExW(&info))
             error::raise("ShellExecuteExW");
     }
+
+    void runas_self(
+        HWND hwnd = NULL,
+        int nShow = SW_NORMAL)
+    {
+        runas(CommandLine::query(), hwnd, nShow);
+    }
 }
diff --git a/src/string.hpp b/src/string.hpp
new file mode 100644
index 0000000..3621ca6
--- /dev/null
+++ b/src/string.hpp
@@ -0,0 +1,154 @@
+#pragma once
+
+#include <cctype>
+#include <cstddef>
+
+#include <algorithm>
+#include <locale>
+#include <sstream>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+namespace string
+{
+    template <typename Char>
+    inline void ltrim(std::basic_string<Char>& s)
+    {
+        s.erase(s.begin(), std::find_if(s.begin(), s.end(), [] (const Char& c)
+        {
+            return !std::isspace(c);
+        }));
+    }
+
+    template <typename Char>
+    inline void rtrim(std::basic_string<Char>& s)
+    {
+        s.erase(std::find_if(s.rbegin(), s.rend(), [] (const Char& c)
+        {
+            return !std::isspace(c);
+        }).base(), s.end());
+    }
+
+    template <typename Char>
+    inline void trim(std::basic_string<Char>& s)
+    {
+        ltrim(s);
+        rtrim(s);
+    }
+
+    template <typename Char, typename Sep, typename InputIterator>
+    inline std::basic_string<Char> join(
+        const Sep& sep,
+        InputIterator beg,
+        InputIterator end)
+    {
+        std::basic_ostringstream<Char> oss;
+
+        if (beg != end)
+        {
+            oss << *beg;
+            ++beg;
+        }
+
+        for (; beg != end; ++beg)
+            oss << sep << *beg;
+
+        return oss.str();
+    }
+
+    template <typename Char, typename Sep>
+    inline std::basic_string<Char> join(
+        const Sep& sep,
+        const std::vector<std::basic_string<Char>>& args)
+    {
+        return join<Char>(sep, args.cbegin(), args.cend());
+    }
+
+    template <typename Char, typename String, typename = void>
+    struct StringHelper
+    { };
+
+    template <typename Char, typename String>
+    struct StringHelper<Char, String, typename std::enable_if<std::is_same<typename std::decay<typename std::remove_pointer<String>::type>::type, Char>::value>::type>
+    {
+        inline StringHelper(const Char& c)
+            : buf{&c}
+            , len{1}
+        { }
+
+        inline StringHelper(const Char* s)
+            : buf{s}
+            , len{std::char_traits<Char>::length(s)}
+        { }
+
+        inline const Char* buffer() const { return buf; }
+
+        inline std::size_t length() const { return len; }
+
+    private:
+        const Char* const buf;
+        const std::size_t len;
+    };
+
+    template <typename Char, typename String>
+    struct StringHelper<Char, String, typename std::enable_if<std::is_same<String, std::basic_string<Char>>::value>::type>
+    {
+        inline StringHelper(const std::basic_string<Char>& s)
+            : s{s}
+        { }
+
+        inline const Char* buffer() const { return s.c_str(); }
+
+        inline std::size_t length() const { return s.length(); }
+
+    private:
+        const std::basic_string<Char>& s;
+    };
+
+    template <typename Char, typename What, typename By>
+    inline void replace(
+        std::basic_string<Char>& s,
+        const What& what,
+        const By& by)
+    {
+        std::size_t pos = 0;
+
+        const StringHelper<Char, typename std::decay<What>::type> what_helper{what};
+        const StringHelper<Char, typename std::decay<By>::type> by_helper{by};
+
+        const auto what_buf = what_helper.buffer();
+        const auto what_len = what_helper.length();
+
+        const auto by_buf = by_helper.buffer();
+        const auto by_len = by_helper.length();
+
+        while ((pos = s.find(what_buf, pos, what_len)) != std::basic_string<Char>::npos)
+        {
+            s.replace(pos, what_len, by_buf, by_len);
+            pos += by_len;
+        }
+    }
+
+    template <typename Char, typename What>
+    inline void prefix_with(
+        std::basic_string<Char>& s,
+        const What& what,
+        const Char& by)
+    {
+        const StringHelper<Char, typename std::decay<What>::type> what_helper{what};
+
+        const auto what_buf = what_helper.buffer();
+        const auto what_len = what_helper.length();
+
+        std::size_t numof_by = 0;
+
+        for (std::size_t i = 0; i < what_len; ++i)
+            numof_by += std::count(s.cbegin(), s.cend(), what_buf[i]);
+
+        s.reserve(s.capacity() + numof_by);
+
+        for (std::size_t i = 0; i < what_len; ++i)
+            replace(s, what_buf[i], by);
+    }
+}