/*

 * MIT License

 *

 * Copyright (c) 2024 Roland Mainz <roland.mainz@nrubsig.org>

 *

 * Permission is hereby granted, free of charge, to any person obtaining a copy

 * of this software and associated documentation files (the "Software"), to deal

 * in the Software without restriction, including without limitation the rights

 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

 * copies of the Software, and to permit persons to whom the Software is

 * furnished to do so, subject to the following conditions:

 *

 * The above copyright notice and this permission notice shall be included in all

 * copies or substantial portions of the Software.

 *

 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE

 * SOFTWARE.

 */

/*

 *

 * Written by Roland Mainz <roland.mainz@nrubsig.org>

 */

#include <windows.h>

#include <stdio.h>

#include <stdbool.h>

#include <Lmcons.h>

#include <process.h>

/*

 * Performance hack:

 * GETTOKINFO_EXTRA_BUFFER - extra space for more data

 * |GetTokenInformation()| for |TOKEN_USER| and |TOKEN_PRIMARY_GROUP|

 * always fails in Win10 with |ERROR_INSUFFICIENT_BUFFER| if you

 * just pass the |sizeof(TOKEN_*)| value. Instead of calling

 * |GetTokenInformation()| with |NULL| arg to obtain the size to

 * allocate we just provide 2048 bytes of extra space after the

 * |TOKEN_*| size, and pray it is enough

 */

#define GETTOKINFO_EXTRA_BUFFER (2048)

bool get_token_primarygroup_name(HANDLE tok, char *out_buffer)

{

    DWORD tokdatalen;

    PTOKEN_PRIMARY_GROUP ptpgroup;

    PSID pgsid;

    DWORD namesize = GNLEN+1;

    char domainbuffer[UNLEN+1];

    DWORD domainbuffer_size = sizeof(domainbuffer);

    SID_NAME_USE name_use;

    tokdatalen = sizeof(TOKEN_PRIMARY_GROUP)+GETTOKINFO_EXTRA_BUFFER;

    ptpgroup = _alloca(tokdatalen);

    if (!GetTokenInformation(tok, TokenPrimaryGroup, ptpgroup,

        tokdatalen, &tokdatalen)) {

            "GetTokenInformation(tok=0x%p, TokenPrimaryGroup) failed, "

            "status=%d\n",

        return false;

    }

    pgsid = ptpgroup->PrimaryGroup;

    if (!LookupAccountSidA(NULL, pgsid, out_buffer, &namesize,

        domainbuffer, &domainbuffer_size, &name_use)) {

            "LookupAccountSidA() failed, status=%d\n",

        return false;

    }

    return true;

}

bool is_group_in_token(HANDLE tok, PSID qsid)

{

    DWORD tokdatalen;

    PTOKEN_GROUPS ptgroups;

    tokdatalen = sizeof(TOKEN_GROUPS)+GETTOKINFO_EXTRA_BUFFER;

    ptgroups = _alloca(tokdatalen);

    if (!GetTokenInformation(tok, TokenGroups, ptgroups,

        tokdatalen, &tokdatalen)) {

            "GetTokenInformation(tok=0x%p, TokenGroups) failed, "

            "status=%d\n",

        return false;

    }

    int i;

    for (i = 0 ; i < ptgroups->GroupCount ; i++) {

        if (EqualSid(qsid, ptgroups->Groups[i].Sid) &&

            (ptgroups->Groups[i].Attributes & SE_GROUP_ENABLED)) {

            return true;

        }

    }

    return false;

}

bool set_token_primarygroup_name(HANDLE tok, const char *new_groupname)

{

    DWORD tokdatalen;

    TOKEN_PRIMARY_GROUP tpgroup;

    char sidbuff[SECURITY_MAX_SID_SIZE+1];

    PSID pgsid = (PSID)sidbuff;

    DWORD pgsid_size = sizeof(sidbuff);

    char domainbuffer[UNLEN+1];

    DWORD domainbuffer_size = sizeof(domainbuffer);

    SID_NAME_USE name_use;

    if (!LookupAccountNameA(NULL, new_groupname,

        pgsid, &pgsid_size, domainbuffer, &domainbuffer_size, &name_use)) {

        return false;

    }

    if (!is_group_in_token(tok, pgsid)) {

        return false;

    }