// Copyright (c) 2016 Egor Tensin <Egor.Tensin@gmail.com>
// This file is part of the "Privilege check" project.
// For details, see https://github.com/egor-tensin/privilege-check.
// Distributed under the MIT License.

#include "error.hpp"
#include "os.hpp"
#include "process.hpp"
#include "resource.hpp"
#include "resource_ids.h"
#include "sid.hpp"
#include "token.hpp"

#include <commctrl.h>
#include <windows.h>
#include <windowsx.h>

#include <exception>
#include <string>

#if defined(_MSC_VER)
#pragma comment(linker, "/manifestdependency:\"type='win32' name='Microsoft.Windows.Common-Controls' version='6.0.0.0' processorArchitecture='*' publicKeyToken='6595b64144ccf1df' language='*'\"")
#endif

namespace
{
    HWND get_dialog_item(HWND wnd, int id)
    {
        const auto lbl = GetDlgItem(wnd, id);
        if (lbl == NULL)
            error::raise("GetDlgItem");
        return lbl;
    }

    bool is_administrator()
    {
        auto token = token::open_for_current_process(
            token::permissions::query() | token::permissions::duplicate());
        if (!token::get_linked(token))
            token = token::get_for_identification(token);

        return token::check_belongs(token, sid::builtin_administrators());
    }

    bool is_run_as_administrator()
    {
        return token::check_belongs(token::dumb(), sid::builtin_administrators());
    }

    bool is_elevated()
    {
        return token::query_elevation(token::open_for_current_process());
    }

    void decorate_elevate_button(HWND wnd, bool decorate)
    {
        const auto btn = get_dialog_item(wnd, IDC_BUTTON_ELEVATE);
        Button_SetElevationRequiredState(btn, !decorate);
    }

    DWORD get_integrity_level()
    {
        return token::query_integrity_level(token::open_for_current_process());
    }

    std::wstring get_integrity_level_as_string()
    {
        return token::integrity_level_to_string(get_integrity_level());
    }

    HWND get_label(HWND wnd, int id)
    {
        return get_dialog_item(wnd, id);
    }

    void set_label(HWND wnd, int id, bool val)
    {
        SetWindowTextW(get_label(wnd, id), val ? L"True" : L"False");
    }

    void set_label(HWND wnd, int id, const wchar_t* s)
    {
        SetWindowTextW(get_label(wnd, id), s);
    }

    void set_label(HWND wnd, int id, const std::wstring& s)
    {
        SetWindowTextW(get_label(wnd, id), s.c_str());
    }

    BOOL on_init_dialog(HWND wnd, HWND, LPARAM)
    {
        try
        {
            set_label(wnd, IDC_ADMINISTRATOR, is_administrator());
        }
        catch (const Error& e)
        {
            set_label(wnd, IDC_ADMINISTRATOR, L"N/A");
            error::report(e);
        }

        try
        {
            set_label(wnd, IDC_RUN_AS_ADMINISTRATOR, is_run_as_administrator());
        }
        catch (const Error& e)
        {
            set_label(wnd, IDC_RUN_AS_ADMINISTRATOR, L"N/A");
            error::report(e);
        }

        if (os::is_vista_or_later())
        {
            try
            {
                const auto elevated = is_elevated();
                set_label(wnd, IDC_ELEVATED, elevated);
                decorate_elevate_button(wnd, elevated);
            }
            catch (const Error& e)
            {
                set_label(wnd, IDC_ELEVATED, L"N/A");
                error::report(e);
            }

            try
            {
                set_label(wnd, IDC_INTEGRITY_LEVEL, get_integrity_level_as_string());
            }
            catch (const Error& e)
            {
                set_label(wnd, IDC_INTEGRITY_LEVEL, L"N/A");
                error::report(e);
            }
        }
        else
        {
            set_label(wnd, IDC_ELEVATED, L"N/A");
            set_label(wnd, IDC_INTEGRITY_LEVEL, L"N/A");
        }

        return TRUE;
    }

    void report_already_elevated(HWND wnd)
    {
        MessageBoxW(wnd,
            resource::load_string(IDS_ALREADY_ELEVATED_TEXT).c_str(),
            resource::load_string(IDS_ALREADY_ELEVATED_CAPTION).c_str(),
            MB_OK);
    }

    void on_button_elevate_click(HWND wnd)
    {
        bool as_admin = false;

        try
        {
            as_admin = is_run_as_administrator();
        }
        catch (const Error& e)
        {
            error::report(e);
            return;
        }

        if (as_admin)
        {
            report_already_elevated(wnd);
            return;
        }

        try
        {
            process::runas_self(wnd);
        }
        catch (const Error& e)
        {
            if (error::get_code(e) != ERROR_CANCELLED)
                error::report(e);
            return;
        }

        EndDialog(wnd, 1);
    }

    void on_command(HWND wnd, int id, HWND, unsigned int)
    {
        switch (id)
        {
            case IDC_BUTTON_ELEVATE:
                on_button_elevate_click(wnd);
                break;

            case IDOK:
            case IDCANCEL:
                EndDialog(wnd, 0);
                break;
        }
    }

    void on_close(HWND wnd)
    {
        EndDialog(wnd, 0);
    }

    INT_PTR CALLBACK dialog_main(
        HWND wnd,
        UINT msg,
        WPARAM wParam,
        LPARAM lParam)
    {
        switch (msg)
        {
            HANDLE_MSG(wnd, WM_INITDIALOG, on_init_dialog);
            HANDLE_MSG(wnd, WM_COMMAND, on_command);
            HANDLE_MSG(wnd, WM_CLOSE, on_close);

            default:
                return FALSE;
        }
    }
}

int APIENTRY wWinMain(
    HINSTANCE instance,
    HINSTANCE,
    wchar_t*,
    int)
{
    try
    {
        const auto ret = DialogBoxW(
            instance,
            MAKEINTRESOURCE(IDD_MAINDIALOG),
            NULL,
            dialog_main);

        switch (ret)
        {
            case -1:
                error::report(error::make("DialogBoxW"));
            default:
                return static_cast<int>(ret);
        }
    }
    catch (const std::exception& e)
    {
        error::report(e);
        return 1;
    }
}