/* Editor Settings: expandtabs and use 4 spaces for indentation
 * ex: set softtabstop=4 tabstop=8 expandtab shiftwidth=4: *
 * -*- mode: c, c-basic-offset: 4 -*- */

/*
 * Copyright Likewise Software
 * All rights reserved.
 *
 * This library is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation; either version 2.1 of the license, or (at
 * your option) any later version.
 *
 * This library 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 Lesser
 * General Public License for more details.  You should have received a copy
 * of the GNU Lesser General Public License along with this program.  If
 * not, see <http://www.gnu.org/licenses/>.
 *
 * LIKEWISE SOFTWARE MAKES THIS SOFTWARE AVAILABLE UNDER OTHER LICENSING
 * TERMS AS WELL.  IF YOU HAVE ENTERED INTO A SEPARATE LICENSE AGREEMENT
 * WITH LIKEWISE SOFTWARE, THEN YOU MAY ELECT TO USE THE SOFTWARE UNDER THE
 * TERMS OF THAT SOFTWARE LICENSE AGREEMENT INSTEAD OF THE TERMS OF THE GNU
 * LESSER GENERAL PUBLIC LICENSE, NOTWITHSTANDING THE ABOVE NOTICE.  IF YOU
 * HAVE QUESTIONS, OR WISH TO REQUEST A COPY OF THE ALTERNATE LICENSING
 * TERMS OFFERED BY LIKEWISE SOFTWARE, PLEASE CONTACT LIKEWISE SOFTWARE AT
 * license@likewisesoftware.com
 */

#include "includes.h"

typedef struct _ACCESS_TOKEN_REP
{
    ULONG Size;
    PACCESS_TOKEN_SELF_RELATIVE pToken;
} ACCESS_TOKEN_REP, *PACCESS_TOKEN_REP;

static LWMsgTypeSpec gAccessTokenRepType[] =
{
    LWMSG_STRUCT_BEGIN(ACCESS_TOKEN_REP),
    LWMSG_MEMBER_UINT32(ACCESS_TOKEN_REP, Size),
    LWMSG_MEMBER_POINTER(ACCESS_TOKEN_REP, pToken, LWMSG_UINT8(BYTE)),
    LWMSG_ATTR_NOT_NULL,
    LWMSG_ATTR_ENCODING("hex+ascii"),
    LWMSG_ATTR_LENGTH_MEMBER(ACCESS_TOKEN_REP, Size),
    LWMSG_STRUCT_END,
    LWMSG_TYPE_END
};

static
LWMsgStatus
MarshalAccessToken(
    LWMsgDataContext* pContext,
    LWMsgType* pType,
    PVOID pPresentObject,
    PVOID pTransmitObject,
    PVOID pUnused
    )
{
    NTSTATUS status = LWMSG_STATUS_SUCCESS;
    PACCESS_TOKEN pToken = *(PACCESS_TOKEN*) pPresentObject;
    PACCESS_TOKEN_REP pRep = (PACCESS_TOKEN_REP) pTransmitObject;

    if (!pToken)
    {
        status = STATUS_INVALID_MESSAGE;
        BAIL_ON_NT_STATUS(status);
    }

    status = RtlAccessTokenToSelfRelativeAccessToken(
            pToken,
            NULL,
            &pRep->Size);
    BAIL_ON_NT_STATUS(status);

    status = RTL_ALLOCATE(
            &pRep->pToken,
            struct _ACCESS_TOKEN_SELF_RELATIVE,
            pRep->Size);
    BAIL_ON_NT_STATUS(status);

    status = RtlAccessTokenToSelfRelativeAccessToken(
            pToken,
            pRep->pToken,
            &pRep->Size);
    BAIL_ON_NT_STATUS(status);

error:

    if (!NT_SUCCESS(status))
    {
        RTL_FREE(&pRep->pToken);
    }

    return NtIpcNtStatusToLWMsgStatus(status);
}

static
LWMsgStatus
UnmarshalAccessToken(
    LWMsgDataContext* pContext,
    LWMsgType* pType,
    PVOID pTransmitObject,
    PVOID pPresentObject,
    PVOID pUnused
    )
{
    NTSTATUS status = LWMSG_STATUS_SUCCESS;
    PACCESS_TOKEN* ppToken = (PACCESS_TOKEN*) pPresentObject;
    PACCESS_TOKEN_REP pRep = (PACCESS_TOKEN_REP) pTransmitObject;

    status = RtlSelfRelativeAccessTokenToAccessToken(
            pRep->pToken,
            pRep->Size,
            ppToken);
    BAIL_ON_NT_STATUS(status);

error:

    return NtIpcNtStatusToLWMsgStatus(status);
}

static
VOID
DestroyPresentedAccessToken(
    LWMsgDataContext* pContext,
    LWMsgType* pType,
    PVOID pPresentObject,
    PVOID pUnused
    )
{
    PACCESS_TOKEN* ppToken = (PACCESS_TOKEN*) pPresentObject;

    RtlReleaseAccessToken(ppToken);
}

static
VOID
DestroyTransmittedAccessToken(
    LWMsgDataContext* pContext,
    LWMsgType* pType,
    PVOID pTransmitObject,
    PVOID pUnused
    )
{
    PACCESS_TOKEN_REP pRep = (PACCESS_TOKEN_REP) pTransmitObject;

    RTL_FREE(&pRep->pToken);
}

static
LWMsgStatus
PrintTokenInfo(
    PACCESS_TOKEN Token,
    TOKEN_INFORMATION_CLASS TokenInformationClass,
    LWMsgBuffer* pBuffer
    )
{
    NTSTATUS status = STATUS_SUCCESS;
    PBYTE tokenInfo = NULL;
    ULONG size = 0;
    PSTR sidString = NULL;

    (void) RtlQueryAccessTokenInformation(
            Token,
            TokenInformationClass,
            tokenInfo,
            size,
            &size);

    status = RTL_ALLOCATE(&tokenInfo, BYTE, size);
    BAIL_ON_NT_STATUS(status);

    status = RtlQueryAccessTokenInformation(
                    Token,
                    TokenInformationClass,
                    tokenInfo,
                    size,
                    &size);
    BAIL_ON_NT_STATUS(status);

    switch (TokenInformationClass)
    {
        case TokenUser:
        {
            PTOKEN_USER tokenUser = (PTOKEN_USER) tokenInfo;

            status = RtlAllocateCStringFromSid(&sidString, tokenUser->User.Sid);
            BAIL_ON_NT_STATUS(status);

            status = NtIpcLWMsgStatusToNtStatus(
                    lwmsg_buffer_print(pBuffer, "    User = %s\n", sidString));
            BAIL_ON_NT_STATUS(status);
            break;
        }
        case TokenGroups:
        {
            PTOKEN_GROUPS tokenGroups = (PTOKEN_GROUPS) tokenInfo;
            ULONG i = 0;

            status = NtIpcLWMsgStatusToNtStatus(
                    lwmsg_buffer_print(pBuffer, "    GroupCount = %u\n", tokenGroups->GroupCount));
            BAIL_ON_NT_STATUS(status);

            status = NtIpcLWMsgStatusToNtStatus(
                    lwmsg_buffer_print(pBuffer, "    Groups =\n    {\n"));
            BAIL_ON_NT_STATUS(status);

            for (i = 0; i < tokenGroups->GroupCount; i++)
            {
                RTL_FREE(&sidString);

                status = RtlAllocateCStringFromSid(&sidString, tokenGroups->Groups[i].Sid);
                BAIL_ON_NT_STATUS(status);

                status = NtIpcLWMsgStatusToNtStatus(
                        lwmsg_buffer_print(pBuffer, "        (0x%08x, %s)\n", tokenGroups->Groups[i].Attributes, sidString));
                BAIL_ON_NT_STATUS(status);
            }

            status = NtIpcLWMsgStatusToNtStatus(
                    lwmsg_buffer_print(pBuffer, "    }\n"));
            BAIL_ON_NT_STATUS(status);
            break;
        }
    }

error:

    RTL_FREE(&sidString);
    RTL_FREE(&tokenInfo);

    return status;
}

static
LWMsgStatus
PrintAccessToken(
    LWMsgDataContext* pContext,
    LWMsgType* pType,
    PVOID pPresentObject,
    PVOID pUnused,
    LWMsgBuffer* pBuffer
    )
{
    NTSTATUS status = STATUS_SUCCESS;
    PACCESS_TOKEN pToken = *(PACCESS_TOKEN*) pPresentObject;

    status = NtIpcLWMsgStatusToNtStatus(
            lwmsg_buffer_print(pBuffer, "\n{\n"));
    BAIL_ON_NT_STATUS(status);

    status = PrintTokenInfo(pToken, TokenUser, pBuffer);
    BAIL_ON_NT_STATUS(status);

    status = PrintTokenInfo(pToken, TokenGroups, pBuffer);
    BAIL_ON_NT_STATUS(status);

    status = NtIpcLWMsgStatusToNtStatus(
            lwmsg_buffer_print(pBuffer, "}"));
    BAIL_ON_NT_STATUS(status);

error:

    return NtIpcNtStatusToLWMsgStatus(status);
}

static LWMsgTypeClass gAccessTokenClass =
{
    .is_pointer = LWMSG_TRUE,
    .transmit_type = gAccessTokenRepType,
    .marshal = MarshalAccessToken,
    .unmarshal = UnmarshalAccessToken,
    .destroy_presented = DestroyPresentedAccessToken,
    .destroy_transmitted = DestroyTransmittedAccessToken,
    .print = PrintAccessToken
};

static
LWMsgTypeSpec gSecurityTokenType[] =
{
    LWMSG_ENUM_BEGIN(enum _LW_IO_CREDS_TYPE, 1, LWMSG_UNSIGNED),
    LWMSG_ENUM_VALUE(IO_CREDS_TYPE_PLAIN),
    LWMSG_ENUM_VALUE(IO_CREDS_TYPE_KRB5_CCACHE),
    LWMSG_ENUM_VALUE(IO_CREDS_TYPE_KRB5_TGT),
    LWMSG_ENUM_VALUE(IO_CREDS_TYPE_IMPERSONATION),
    LWMSG_ENUM_END,
    LWMSG_TYPE_END
};

static
LWMsgStatus
VerifyInnerCreds(
    LWMsgDataContext* pContext,
    LWMsgBool Unmarshal,
    PVOID pObject,
    PVOID pUnused
    )
{
    PIO_CREDS* ppCreds = (PIO_CREDS*) pObject;

    if (*ppCreds && (*ppCreds)->type == IO_CREDS_TYPE_IMPERSONATION)
    {
        return LWMSG_STATUS_MALFORMED;
    }

    return LWMSG_STATUS_SUCCESS;
}

LWMsgTypeSpec gSecurityTokenRepSpec[] =
{
    /* Begin structure */
    LWMSG_STRUCT_BEGIN(IO_CREDS),
    /* Discriminator */
    LWMSG_MEMBER_TYPESPEC(IO_CREDS, type, gSecurityTokenType),
    /* Begin union */
    LWMSG_MEMBER_UNION_BEGIN(IO_CREDS, payload),
    /* Union arm -- plain */
    LWMSG_MEMBER_STRUCT_BEGIN(union _LW_IO_CREDS_U, plain),
    LWMSG_MEMBER_PWSTR(struct _LW_IO_CREDS_PLAIN, pwszUsername),
    LWMSG_MEMBER_PWSTR(struct _LW_IO_CREDS_PLAIN, pwszDomain),
    LWMSG_MEMBER_PWSTR(struct _LW_IO_CREDS_PLAIN, pwszPassword),
    LWMSG_ATTR_SENSITIVE,
    LWMSG_STRUCT_END,
    LWMSG_ATTR_TAG(IO_CREDS_TYPE_PLAIN),
    /* Union arm -- krb5 ccache */
    LWMSG_MEMBER_STRUCT_BEGIN(union _LW_IO_CREDS_U, krb5Ccache),
    LWMSG_MEMBER_PWSTR(struct _LW_IO_CREDS_KRB5_CCACHE, pwszPrincipal),
    LWMSG_MEMBER_PWSTR(struct _LW_IO_CREDS_KRB5_CCACHE, pwszCachePath),
    LWMSG_STRUCT_END,
    LWMSG_ATTR_TAG(IO_CREDS_TYPE_KRB5_CCACHE),
    /* Union arm -- krb5 tgt */
    LWMSG_MEMBER_STRUCT_BEGIN(union _LW_IO_CREDS_U, krb5Tgt),
    LWMSG_MEMBER_PWSTR(struct _LW_IO_CREDS_KRB5_TGT, pwszClientPrincipal),
    LWMSG_MEMBER_PWSTR(struct _LW_IO_CREDS_KRB5_TGT, pwszServerPrincipal),
    LWMSG_MEMBER_UINT32(struct _LW_IO_CREDS_KRB5_TGT, authTime),
    LWMSG_MEMBER_UINT32(struct _LW_IO_CREDS_KRB5_TGT, startTime),
    LWMSG_MEMBER_UINT32(struct _LW_IO_CREDS_KRB5_TGT, endTime),
    LWMSG_MEMBER_UINT32(struct _LW_IO_CREDS_KRB5_TGT, renewTillTime),
    LWMSG_MEMBER_UINT32(struct _LW_IO_CREDS_KRB5_TGT, keyType),
    LWMSG_MEMBER_UINT32(struct _LW_IO_CREDS_KRB5_TGT, ulKeySize),
    LWMSG_MEMBER_POINTER(struct _LW_IO_CREDS_KRB5_TGT, pKeyData, LWMSG_UINT8(BYTE)),
    LWMSG_ATTR_SENSITIVE,
    LWMSG_ATTR_LENGTH_MEMBER(struct _LW_IO_CREDS_KRB5_TGT, ulKeySize),
    LWMSG_MEMBER_UINT32(struct _LW_IO_CREDS_KRB5_TGT, tgtFlags),
    LWMSG_MEMBER_UINT32(struct _LW_IO_CREDS_KRB5_TGT, ulTgtSize),
    LWMSG_MEMBER_POINTER(struct _LW_IO_CREDS_KRB5_TGT, pTgtData, LWMSG_UINT8(BYTE)),
    LWMSG_ATTR_SENSITIVE,
    LWMSG_ATTR_LENGTH_MEMBER(struct _LW_IO_CREDS_KRB5_TGT, ulTgtSize),
    LWMSG_STRUCT_END,
    LWMSG_ATTR_TAG(IO_CREDS_TYPE_KRB5_TGT),
    LWMSG_MEMBER_STRUCT_BEGIN(union _LW_IO_CREDS_U, Impersonation),
    LWMSG_MEMBER_CUSTOM(struct _LW_IO_CREDS_IMPERSONATION, pToken, &gAccessTokenClass, NULL),
    LWMSG_MEMBER_POINTER(struct _LW_IO_CREDS_IMPERSONATION, pCreds, LWMSG_TYPESPEC(gSecurityTokenRepSpec)),
    LWMSG_ATTR_VERIFY(VerifyInnerCreds, NULL),
    LWMSG_STRUCT_END,
    LWMSG_ATTR_TAG(IO_CREDS_TYPE_IMPERSONATION),
    /* End union */
    LWMSG_UNION_END,
    LWMSG_ATTR_DISCRIM(IO_CREDS, type),
    /* End structure */
    LWMSG_STRUCT_END,
    LWMSG_TYPE_END
};

static LWMsgTypeSpec gStatusReplySpec[] =
{
    /* Begin structure */
    LWMSG_STRUCT_BEGIN(LWIO_STATUS_REPLY),
    /* err - marshal as 32-bit unsigned integer */
    LWMSG_MEMBER_UINT32(LWIO_STATUS_REPLY, dwError),
    /* End structure */
    LWMSG_STRUCT_END,
    LWMSG_TYPE_END
};

static LWMsgTypeSpec gDriverNameSpec[] =
{
    LWMSG_PWSTR,
    LWMSG_ATTR_NOT_NULL,
    LWMSG_TYPE_END
};

static LWMsgTypeSpec gDriverStateSpec[] =
{
    LWMSG_UINT8(LWIO_DRIVER_STATE),
    LWMSG_ATTR_RANGE(LWIO_DRIVER_STATE_UNKNOWN, LWIO_DRIVER_STATE_LOADED),
    LWMSG_TYPE_END
};

static LWMsgProtocolSpec gLwIoDaemonProtocolSpec[] =
{
    LWMSG_MESSAGE(LWIO_QUERY_STATE_DRIVER,          gDriverNameSpec),
    LWMSG_MESSAGE(LWIO_QUERY_STATE_DRIVER_SUCCESS,  gDriverStateSpec),
    LWMSG_MESSAGE(LWIO_QUERY_STATE_DRIVER_FAILED,   gStatusReplySpec),
    LWMSG_MESSAGE(LWIO_LOAD_DRIVER,                 gDriverNameSpec),
    LWMSG_MESSAGE(LWIO_LOAD_DRIVER_SUCCESS,         gStatusReplySpec),
    LWMSG_MESSAGE(LWIO_LOAD_DRIVER_FAILED,          gStatusReplySpec),
    LWMSG_MESSAGE(LWIO_UNLOAD_DRIVER,               gDriverNameSpec),
    LWMSG_MESSAGE(LWIO_UNLOAD_DRIVER_SUCCESS,       gStatusReplySpec),
    LWMSG_MESSAGE(LWIO_UNLOAD_DRIVER_FAILED,        gStatusReplySpec),
    LWMSG_PROTOCOL_END
};

NTSTATUS
LwIoDaemonIpcAddProtocolSpec(
    IN OUT LWMsgProtocol* pProtocol
    )
{
    NTSTATUS status = 0;
    int EE = 0;
    LWMsgStatus msgStatus = 0;

    msgStatus = lwmsg_protocol_add_protocol_spec(pProtocol, gLwIoDaemonProtocolSpec);
    status = NtIpcLWMsgStatusToNtStatus(msgStatus);
    GOTO_CLEANUP_ON_STATUS_EE(status, EE);

cleanup:

    LOG_LEAVE_IF_STATUS_EE(status, EE);

    return status;
}
