// Copyright (C) Microsoft.  All rights reserved.

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <objbase.h>
#include <netcon.h>
#include <stdio.h>

#pragma comment(lib, "ole32.lib")
#pragma comment(lib, "oleaut32.lib")

// as in winsock.h
#define NAT_PROTOCOL_TCP 6

HRESULT DeletePortMapping(INetSharingManager *pNSM,
                          UCHAR               ucIPProtocol,
                          short               usExternalPort) {    // this is done in 2 parts:
    // 1:  enum connections until we get one that we can convert into an INetSharingConfiguration
    // 2:  then, enum portmappings, and delete if we find a match.

    // PART 1: find a valid connection
    INetConnection *pNC = nullptr;    // fill this out for part 2 below

    INetSharingEveryConnectionCollection *pNSECC = nullptr;
    HRESULT                               hr     = pNSM->get_EnumEveryConnection(&pNSECC);
    if (!pNSECC) {
        wprintf(L"failed to get EveryConnectionCollection!\r\n");
    } else {
        // enumerate connections
        IEnumVARIANT *pEV  = nullptr;
        IUnknown     *pUnk = nullptr;
        hr                 = pNSECC->get__NewEnum(&pUnk);
        if (pUnk) {
            hr = pUnk->QueryInterface(__uuidof(IEnumVARIANT), (void **)&pEV);
            pUnk->Release();
        }
        if (pEV) {
            VARIANT v;
            VariantInit(&v);
            BOOL bFoundIt = FALSE;
            while (S_OK == pEV->Next(1, &v, nullptr)) {
                if (V_VT(&v) == VT_UNKNOWN) {
                    V_UNKNOWN(&v)->QueryInterface(__uuidof(INetConnection), (void **)&pNC);
                    if (pNC) {
                        INetConnectionProps *pNCP = nullptr;
                        pNSM->get_NetConnectionProps(pNC, &pNCP);
                        if (!pNCP) {
                            wprintf(L"failed to get NetConnectionProps!\r\n");
                        } else {
                            // check properties for firewalled or shared connection
                            DWORD dwCharacteristics = 0;
                            pNCP->get_Characteristics(&dwCharacteristics);
                            if (dwCharacteristics & (NCCF_SHARED | NCCF_FIREWALLED)) {
                                NETCON_MEDIATYPE MediaType = NCM_NONE;
                                pNCP->get_MediaType(&MediaType);
                                if ((MediaType != NCM_SHAREDACCESSHOST_LAN) &&(MediaType != NCM_SHAREDACCESSHOST_RAS)) {
                                    // got a shared/firewalled connection
                                    bFoundIt = TRUE;
                                }
                            }
                            pNCP->Release();
                        }
                        if (bFoundIt == FALSE) {
                            pNC->Release();
                            pNC = nullptr;
                        }
                    }
                }
                VariantClear(&v);
                if (bFoundIt == TRUE) {
                    break;
                }
            }
            pEV->Release();
        }
        pNSECC->Release();
    }
    if (pNC == nullptr) {
        wprintf(L"failed to find a valid connection!\r\n");
        return E_FAIL;
    }
    INetSharingConfiguration *pNSC = nullptr;
    hr                             = pNSM->get_INetSharingConfigurationForINetConnection(pNC, &pNSC);
    pNC->Release();    // don't need this anymore
    if (!pNSC) {
        wprintf(L"can't make INetSharingConfiguration object!\r\n");
        return hr;
    }

    // PART 2:  enum port mappings, deleting match, if any
    INetSharingPortMappingCollection *pNSPMC = nullptr;
    hr                                       = pNSC->get_EnumPortMappings(ICSSC_DEFAULT, &pNSPMC);
    if (!pNSPMC) {
        wprintf(L"can't get PortMapping collection!\r\n");
    } else {

        // this is the interface to be filled out by the code below
        INetSharingPortMapping *pNSPM = nullptr;

        IEnumVARIANT *pEV  = nullptr;
        IUnknown     *pUnk = nullptr;
        hr                 = pNSPMC->get__NewEnum(&pUnk);
        if (pUnk) {
            hr = pUnk->QueryInterface(__uuidof(IEnumVARIANT), (void **)&pEV);
            pUnk->Release();
        }
        if (pEV) {
            VARIANT v;
            VariantInit(&v);
            BOOL bFoundIt = FALSE;
            while (S_OK == pEV->Next(1, &v, nullptr)) {
                if (V_VT(&v) == VT_DISPATCH) {
                    V_DISPATCH(&v)->QueryInterface(__uuidof(INetSharingPortMapping), (void **)&pNSPM);
                    if (pNSPM) {
                        INetSharingPortMappingProps *pNSPMP = nullptr;
                        hr                                  = pNSPM->get_Properties(&pNSPMP);
                        if (pNSPMP) {
                            UCHAR uc = 0;
                            pNSPMP->get_IPProtocol(&uc);
                            long usExternal = 0;
                            pNSPMP->get_ExternalPort(&usExternal);
                            if ((uc == ucIPProtocol) &&(usExternal == usExternalPort)) {
                                bFoundIt = TRUE;
                            }

                            pNSPMP->Release();
                        }
                        if (bFoundIt == FALSE) {    // hang onto reference to pNSPM iff found (used below)
                            pNSPM->Release();
                            pNSPM = nullptr;
                        }
                    }
                }
                VariantClear(&v);
                if (bFoundIt == TRUE) {
                    break;    // bail out if we've found one
                }
            }
            pEV->Release();
        }

        if (pNSPM) {
            hr = pNSPM->Delete();    // or pNSC->RemovePortMapping (pNSPM);
            wprintf(L"just deleted a port mapping!\r\n");

            pNSPM->Release();
        }
        pNSPMC->Release();
    }
    pNSC->Release();
    return hr;
}

HRESULT AddAsymmetricPortMapping(INetSharingConfiguration *pNSC) {
    HRESULT      hr  = S_OK;
#if 0    
    VARIANT_BOOL vb1 = VARIANT_FALSE;
    VARIANT_BOOL vb2 = VARIANT_FALSE;
    pNSC->get_SharingEnabled(&vb1);
    pNSC->get_InternetFirewallEnabled(&vb2);
    if ((vb1 == VARIANT_FALSE) &&(vb2 == VARIANT_FALSE)) {
        wprintf(L"sharing and/or firewall not enabled on this connection!\r\n");
    } else {
        INetSharingPortMapping *pNSPM = nullptr;
        hr                            = pNSC->AddPortMapping("Ben's Port Mapping",
                                  NAT_PROTOCOL_TCP,
                                  555,
                                  444,
                                  0,
                                  "192.168.0.2",
                                  ICSTT_IPADDRESS,
                                  &pNSPM);
        if (pNSPM) {
            wprintf(L"just added NAT_PROTOCOL_TCP, 555, 444!\r\n");

            hr = pNSPM->Enable();
            wprintf(L"just enabled port mapping!\r\n");

            pNSPM->Release();
        } else {
            wprintf(L"failed to add asymmetric port mapping!\r\n");
        }
    }
#endif
    return hr;
}

HRESULT DoTheWork(INetSharingManager *pNSM) {    // add a port mapping to every firewalled or shared connection

    INetSharingEveryConnectionCollection *pNSECC = nullptr;
    HRESULT                               hr     = pNSM->get_EnumEveryConnection(&pNSECC);
    if (!pNSECC) {
        wprintf(L"failed to get EveryConnectionCollection!\r\n");
    } else {

        // enumerate connections
        IEnumVARIANT *pEV  = nullptr;
        IUnknown     *pUnk = nullptr;
        hr                 = pNSECC->get__NewEnum(&pUnk);
        if (pUnk) {
            hr = pUnk->QueryInterface(__uuidof(IEnumVARIANT), (void **)&pEV);
            pUnk->Release();
        }
        if (pEV) {
            VARIANT v;
            VariantInit(&v);
            while (S_OK == pEV->Next(1, &v, nullptr)) {
                if (V_VT(&v) == VT_UNKNOWN) {
                    INetConnection *pNC = nullptr;
                    V_UNKNOWN(&v)->QueryInterface(__uuidof(INetConnection), (void **)&pNC);
                    if (pNC) {
                        INetConnectionProps *pNCP = nullptr;
                        pNSM->get_NetConnectionProps(pNC, &pNCP);
                        if (!pNCP) {
                            wprintf(L"failed to get NetConnectionProps!\r\n");
                        } else {
                            // check properties for firewalled or shared connection
                            DWORD dwCharacteristics = 0;
                            pNCP->get_Characteristics(&dwCharacteristics);
                            if (dwCharacteristics & (NCCF_SHARED | NCCF_FIREWALLED)) {
                                NETCON_MEDIATYPE MediaType = NCM_NONE;
                                pNCP->get_MediaType(&MediaType);
                                if ((MediaType != NCM_SHAREDACCESSHOST_LAN) &&
                                    (MediaType != NCM_SHAREDACCESSHOST_RAS)) {
                                    // got a shared/firewalled connection
                                    INetSharingConfiguration *pNSC = nullptr;
                                    hr = pNSM->get_INetSharingConfigurationForINetConnection(pNC, &pNSC);
                                    if (!pNSC) {
                                        wprintf(L"can't make INetSharingConfiguration object!\r\n");
                                    } else {
                                        hr = AddAsymmetricPortMapping(pNSC);
                                        pNSC->Release();
                                    }
                                }
                            }
                            pNCP->Release();
                        }
                        pNC->Release();
                    }
                }
                VariantClear(&v);
            }
            pEV->Release();
        }
        pNSECC->Release();
    }
    return hr;
}

int main3() {
    CoInitialize(nullptr);

    // init security to enum RAS connections
    CoInitializeSecurity(nullptr,
                         -1,
                         nullptr,
                         nullptr,
                         RPC_C_AUTHN_LEVEL_PKT,
                         RPC_C_IMP_LEVEL_IMPERSONATE,
                         nullptr,
                         EOAC_NONE,
                         nullptr);

    INetSharingManager *pNSM = nullptr;
    HRESULT             hr   = ::CoCreateInstance(CLSID_NetSharingManager,
                                    nullptr,
                                    CLSCTX_ALL,
                                    IID_INetSharingManager,
                                    reinterpret_cast<void **>(&pNSM));
    if (!pNSM) {
        wprintf(L"failed to create NetSharingManager object\r\n");
    } else {
        // in case it exists already
        DeletePortMapping(pNSM, NAT_PROTOCOL_TCP, 555);

        // add a port mapping to every shared or firewalled connection.
        hr = DoTheWork(pNSM);
        if (SUCCEEDED(hr)) {

            // do other work here.
            // when you're done,
            // clean up port mapping

            hr = DeletePortMapping(pNSM, NAT_PROTOCOL_TCP, 555);
        }

        pNSM->Release();
    }

    CoUninitialize();
    return (int)hr;
}