/********************************************************************************
 Copyright (C) 2012 Hugh Bailey <obs.jim@gmail.com>
                    Richard Stanway

 This program is free software; you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published by
 the Free Software Foundation; either version 2 of the License, or
 (at your option) any later version.

 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with this program; if not, write to the Free Software
 Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA.
********************************************************************************/

#include "Main.h"
#include <winhttp.h>

__declspec(thread) bool invalidCN = false;

static void CALLBACK WinHTTPStatusCallback(HINTERNET hInternet, DWORD_PTR dwContext, DWORD dwInternetStatus, LPVOID lpvStatusInformation, DWORD dwStatusInformationLength)
{
    if (dwStatusInformationLength != 4)
        return;

    if ((*(DWORD*)lpvStatusInformation) & ~(WINHTTP_CALLBACK_STATUS_FLAG_CERT_CN_INVALID | WINHTTP_CALLBACK_STATUS_FLAG_INVALID_CA))
        return;

    if ((*(DWORD*)lpvStatusInformation) & WINHTTP_CALLBACK_STATUS_FLAG_CERT_CN_INVALID)
        invalidCN = true;
}

BOOL HTTPGetFile(CTSTR url, CTSTR outputPath, CTSTR extraHeaders, int *responseCode, HTTPGetFileWin81TLSSNIBugHandler h)
{
    HINTERNET hSession = NULL;
    HINTERNET hConnect = NULL;
    HINTERNET hRequest = NULL;
    URL_COMPONENTS  urlComponents;
    BOOL secure = FALSE;
    BOOL ret = FALSE;

    String hostName, path;

    const TCHAR *acceptTypes[] = {
        TEXT("*/*"),
        NULL
    };

    hostName.SetLength(256);
    path.SetLength(1024);

    zero(&urlComponents, sizeof(urlComponents));

    urlComponents.dwStructSize = sizeof(urlComponents);
    
    urlComponents.lpszHostName = hostName;
    urlComponents.dwHostNameLength = hostName.Length();

    urlComponents.lpszUrlPath = path;
    urlComponents.dwUrlPathLength = path.Length();

    WinHttpCrackUrl(url, 0, 0, &urlComponents);

    if (urlComponents.nPort == 443)
        secure = TRUE;

retry:
    hSession = WinHttpOpen(OBS_VERSION_STRING, WINHTTP_ACCESS_TYPE_DEFAULT_PROXY, WINHTTP_NO_PROXY_NAME, WINHTTP_NO_PROXY_BYPASS, 0);
    if (!hSession)
        goto failure;

    if (secure)
        WinHttpSetStatusCallback(hSession, WinHTTPStatusCallback, WINHTTP_CALLBACK_FLAG_SECURE_FAILURE, NULL);

    hConnect = WinHttpConnect(hSession, hostName, secure ? INTERNET_DEFAULT_HTTPS_PORT : INTERNET_DEFAULT_HTTP_PORT, 0);
    if (!hConnect)
        goto failure;

    hRequest = WinHttpOpenRequest(hConnect, TEXT("GET"), path, NULL, WINHTTP_NO_REFERER, acceptTypes, secure ? WINHTTP_FLAG_SECURE|WINHTTP_FLAG_REFRESH : WINHTTP_FLAG_REFRESH);
    if (!hRequest)
        goto failure;

    BOOL bResults = WinHttpSendRequest(hRequest, extraHeaders, extraHeaders ? -1 : 0, WINHTTP_NO_REQUEST_DATA, 0, 0, 0);

    if (!bResults)
    {
        DWORD err = GetLastError();
        if (err == ERROR_WINHTTP_SECURE_FAILURE && invalidCN && h && h())
        {
            secure = 0;
            WinHttpSetStatusCallback(hSession, nullptr, WINHTTP_CALLBACK_FLAG_SECURE_FAILURE, NULL);
            goto retry;
        }
        goto failure;
    }

    // End the request.
    bResults = WinHttpReceiveResponse(hRequest, NULL);

    TCHAR statusCode[8];
    DWORD statusCodeLen;

    statusCodeLen = sizeof(statusCode);
    if (!WinHttpQueryHeaders (hRequest, WINHTTP_QUERY_STATUS_CODE, WINHTTP_HEADER_NAME_BY_INDEX, &statusCode, &statusCodeLen, WINHTTP_NO_HEADER_INDEX))
        goto failure;

    *responseCode = wcstoul(statusCode, NULL, 10);

    if (bResults && *responseCode == 200)
    {
        BYTE buffer[16384];
        DWORD dwSize, dwOutSize;

        XFile updateFile;

        if (!updateFile.Open(outputPath, XFILE_WRITE, CREATE_ALWAYS))
            goto failure;

        do 
        {
            // Check for available data.
            dwSize = 0;
            if (!WinHttpQueryDataAvailable(hRequest, &dwSize))
                goto failure;

            if (!WinHttpReadData(hRequest, (LPVOID)buffer, dwSize, &dwOutSize))
            {
                goto failure;
            }
            else
            {
                if (!dwOutSize)
                    break;

                if (!updateFile.Write(buffer, dwOutSize))
                    goto failure;
            }
        } while (dwSize > 0);

        updateFile.Close();
    }

    ret = TRUE;

failure:
    if (hSession)
    {
        WinHttpSetStatusCallback(hSession, nullptr, WINHTTP_CALLBACK_FLAG_SECURE_FAILURE, NULL);
        WinHttpCloseHandle(hSession);
    }
    if (hConnect)
        WinHttpCloseHandle(hConnect);
    if (hRequest)
        WinHttpCloseHandle(hRequest);

    return ret;
}

String CreateHTTPURL(String host, String path, String extra, bool secure)
{
    URL_COMPONENTS components = {
        sizeof URL_COMPONENTS,
        secure ? L"https" : L"http",
        secure ? 5 : 4,
        secure ? INTERNET_SCHEME_HTTPS : INTERNET_SCHEME_HTTP,
        host.Array(), host.Length(),
        secure ? INTERNET_DEFAULT_HTTPS_PORT : INTERNET_DEFAULT_HTTP_PORT,
        nullptr, 0,
        nullptr, 0,
        path.Array(), path.Length(),
        extra.Array(), extra.Length()
    };

    String url;
    url.SetLength(MAX_PATH);
    DWORD length = MAX_PATH;
    if (!WinHttpCreateUrl(&components, ICU_ESCAPE, url.Array(), &length))
        return String();

    url.SetLength(length);
    return url;
}