diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/cmd_line.hpp | 194 | ||||
| -rw-r--r-- | src/main.cpp | 2 | ||||
| -rw-r--r-- | src/process.hpp | 27 | ||||
| -rw-r--r-- | src/string.hpp | 154 |
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); + } +} |