/*

 *

 *    Copyright (c) 2013-2017 Nest Labs, Inc.

 *    All rights reserved.

 *

 *    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

 *

 *        http://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.

 */

/**

 *    @file

 *      A library for manipulating PASE state for testing and fuzzing

 *

 */

#include <stdio.h>

#include <vector>

#include "ToolCommon.h"

#include "PASEEngineTest.h"

#include <Weave/Support/ErrorStr.h>

#include <Weave/Profiles/security/WeaveSecurity.h>

#include <Weave/Profiles/security/WeavePASE.h>

#include <Weave/Support/RandUtils.h>

#if WEAVE_SYSTEM_CONFIG_USE_LWIP

#include "lwip/tcpip.h"

#endif // WEAVE_SYSTEM_CONFIG_USE_LWIP

using namespace nl::Weave::Profiles::Security;

using namespace nl::Weave::Profiles::Security::PASE;

using System::PacketBuffer;

#define TOOL_NAME "TestPASE"

const char * INITIATOR_STEP_1 = "InitiatorStep1";

const char * RESPONDER_RECONFIGURE = "ResponderReconfigure";

const char * RESPONDER_STEP_1 = "ResponderStep1";

const char * RESPONDER_STEP_2 = "ResponderStep2";

const char * INITIATOR_STEP_2 = "InitiatorStep2";

const char * RESPONDER_KEY_CONFIRM = "ResponderKeyConfirm";

#define VerifyOrQuit(TST, MSG) \

do { \

    if (!(TST)) \

    { \

        fprintf(stderr, "%s FAILED: ", __FUNCTION__); \

        fputs(MSG, stderr); \

        exit(-1); \

    } \

} while (0)

#define SuccessOrQuit(ERR, MSG) \

do { \

    if ((ERR) != WEAVE_NO_ERROR) \

    { \

        fprintf(stderr, "%s FAILED: ", __FUNCTION__); \

        fputs(MSG, stderr); \

        fputs(ErrorStr(ERR), stderr); \

        fputs("\n", stderr); \

        exit(-1); \

    } \

} while (0)

void MessageMutator::MutateMessage(const char *msgName, PacketBuffer *msgBuf) { }

static MessageMutator gNullMutator;

MessageExternalFuzzer::MessageExternalFuzzer(const char *msgType)

{

    mMsgType = msgType;

    mSaveCorpus = false;

}

void MessageExternalFuzzer::MutateMessage(const char *msgType, PacketBuffer *msgBuf)

{

    if (strcmp(msgType, mMsgType) == 0)

    {

        uint8_t *msgStart = msgBuf->Start();

        if (mSaveCorpus)

        {

            MessageExternalFuzzer::SaveCorpus(msgStart, msgBuf->DataLength(), mMsgType);

        }

        msgBuf->SetDataLength(mFuzzInputSize);

        memcpy(msgStart, mFuzzInput, msgBuf->DataLength());

    }

}

void MessageExternalFuzzer::SaveCorpus(const uint8_t *inBuf, size_t size, const char *fileName)

{

    FILE* file = fopen(fileName, "wb+" );

    VerifyOrQuit(file != NULL, "Could not open file");

    VerifyOrQuit((fwrite(inBuf, 1, size, file) == size), "Could not write corpus file.");

    fclose(file);

}

MessageExternalFuzzer& MessageExternalFuzzer::SaveCorpusFile(bool val) { mSaveCorpus = val; return *this; }

MessageExternalFuzzer& MessageExternalFuzzer::FuzzInput(const uint8_t *val, size_t size) { mFuzzInput = val; mFuzzInputSize = size; return *this; }

//Start PASEEngineTest Initialization

PASEEngineTest::PASEEngineTest(const char *testName)

{

    mTestName = testName;

    mProposedConfig = mExpectedConfig = kPASEConfig_Unspecified;

    mInitPW = mRespPW = "TestPassword";

    mInitiatorAllowedConfigs = mResponderAllowedConfigs = kPASEConfig_Config1|kPASEConfig_Config4;

    mExpectReconfig = false;

    mForceRepeatedReconfig = false;

    memset(mExpectedErrors, 0, sizeof(mExpectedErrors));

    mMutator = &gNullMutator;

    mLogMessageData = false;

}

const char *PASEEngineTest::TestName() const { return mTestName; }

uint32_t PASEEngineTest::ProposedConfig() const { return mProposedConfig; }

PASEEngineTest& PASEEngineTest::ProposedConfig(uint32_t val) { mProposedConfig = val; return *this; }

uint32_t PASEEngineTest::InitiatorAllowedConfigs() const { return mInitiatorAllowedConfigs; }

PASEEngineTest& PASEEngineTest::InitiatorAllowedConfigs(uint32_t val) { mInitiatorAllowedConfigs = val; return *this; }

uint32_t PASEEngineTest::ResponderAllowedConfigs() const { return mResponderAllowedConfigs; }

PASEEngineTest& PASEEngineTest::ResponderAllowedConfigs(uint32_t val) { mResponderAllowedConfigs = val; return *this; }

const char* PASEEngineTest::InitiatorPassword() const { return mInitPW; }

PASEEngineTest& PASEEngineTest::InitiatorPassword(const char* val) { mInitPW = val; return *this; }

const char* PASEEngineTest::ResponderPassword() const { return mRespPW; }

PASEEngineTest& PASEEngineTest::ResponderPassword(const char* val) { mRespPW = val; return *this; }

uint32_t PASEEngineTest::ExpectReconfig() const { return mExpectReconfig; }

PASEEngineTest& PASEEngineTest::ExpectReconfig(uint32_t expectedConfig)

{

    mExpectReconfig = true;

    mExpectedConfig = expectedConfig;

    return *this;

}

uint32_t PASEEngineTest::ExpectedConfig() const { return mExpectedConfig != kPASEConfig_Unspecified ? mExpectedConfig : mProposedConfig; }

bool PASEEngineTest::PerformReconfig() const { return mForceRepeatedReconfig; }

PASEEngineTest& PASEEngineTest::PerformReconfig(bool val) { mForceRepeatedReconfig = val; return *this; }

bool PASEEngineTest::ConfirmKey() const { return mConfirmKey; }

PASEEngineTest& PASEEngineTest::ConfirmKey(bool val) { mConfirmKey = val; return *this; }

PASEEngineTest& PASEEngineTest::ExpectError(WEAVE_ERROR err)

{

    return ExpectError(NULL, err);

}

PASEEngineTest& PASEEngineTest::ExpectError(const char *opName, WEAVE_ERROR err)

{

    for (size_t i = 0; i < kMaxExpectedErrors; i++)

    {

        if (mExpectedErrors[i].Error == WEAVE_NO_ERROR)

        {

            mExpectedErrors[i].Error = err;

            mExpectedErrors[i].OpName = opName;

            break;

        }

    }

    return *this;

}

bool PASEEngineTest::IsExpectedError(const char *opName, WEAVE_ERROR err) const

{

    for (size_t i = 0; i < kMaxExpectedErrors && mExpectedErrors[i].Error != WEAVE_NO_ERROR; i++)

    {

        if (mExpectedErrors[i].Error == err &&

            (mExpectedErrors[i].OpName == NULL || strcmp(mExpectedErrors[i].OpName, opName) == 0))

        {

            return true;

        }

    }

    return false;

}

bool PASEEngineTest::IsSuccessExpected() const { return mExpectedErrors[0].Error == WEAVE_NO_ERROR; }

PASEEngineTest& PASEEngineTest::Mutator(MessageMutator *mutator) { mMutator = mutator; return *this; }

bool PASEEngineTest::LogMessageData() const { return mLogMessageData; }

PASEEngineTest& PASEEngineTest::LogMessageData(bool val) { mLogMessageData = val; return *this; }

//private

void PASEEngineTest::setAllowedResponderConfigs(WeavePASEEngine &responderEng)

{

#if WEAVE_CONFIG_SUPPORT_PASE_CONFIG0_TEST_ONLY

        if (mResponderAllowedConfigs == kPASEConfig_Config0_TEST_ONLY)

            responderEng.AllowedPASEConfigs = kPASEConfig_SupportConfig0Bit_TEST_ONLY;

        else

#endif

#if WEAVE_CONFIG_SUPPORT_PASE_CONFIG1

        if (mResponderAllowedConfigs == kPASEConfig_Config1)

            responderEng.AllowedPASEConfigs = kPASEConfig_SupportConfig1Bit;

        else

#endif

#if WEAVE_CONFIG_SUPPORT_PASE_CONFIG2

        if (mResponderAllowedConfigs == kPASEConfig_Config2)

            respEng.AllowedPASEConfigs = kPASEConfig_SupportConfig2Bit;

        else

#endif

#if WEAVE_CONFIG_SUPPORT_PASE_CONFIG3

        if (mResponderAllowedConfigs == kPASEConfig_Config3)

            responderEng.AllowedPASEConfigs = kPASEConfig_SupportConfig3Bit;

        else

#endif

#if WEAVE_CONFIG_SUPPORT_PASE_CONFIG4

        if (mResponderAllowedConfigs == kPASEConfig_Config4)

            responderEng.AllowedPASEConfigs = kPASEConfig_SupportConfig4Bit;

        else

#endif

#if WEAVE_CONFIG_SUPPORT_PASE_CONFIG5

        if (mResponderAllowedConfigs == kPASEConfig_Config5)

            responderEng.AllowedPASEConfigs = kPASEConfig_SupportConfig5Bit;

        else

#endif

            responderEng.AllowedPASEConfigs = 0x0;

}

enum

{

    kMaxExpectedErrors = 32

};

//end PASEEngineTest Initialization

void PASEEngineTest::Run()

{

    WEAVE_ERROR err;

    WeavePASEEngine initiatorEng;

    WeavePASEEngine responderEng;

    PacketBuffer *msgBuf = NULL;

    PacketBuffer *msgBuf2 = NULL;

    WeaveFabricState initFabricState;

    WeaveFabricState respFabricState;

    const WeaveEncryptionKey *initiatorKey;

    const WeaveEncryptionKey *responderKey;

    uint64_t initNodeId = 1;

    uint64_t respNodeId = 2;

    uint16_t sessionKeyId = sTestDefaultSessionKeyId;

    uint16_t encType = kWeaveEncryptionType_AES128CTRSHA1;

    uint16_t pwSrc = kPasswordSource_PairingCode;

    bool expectSuccess = strcmp(mInitPW, mRespPW) == 0;

    if (LogMessageData())

    {

        printf("========== Starting Test: %s\n", TestName());

        printf("Pr: %d\nex: %d\n", ProposedConfig(), ExpectedConfig());

    }

    initiatorEng.Init();

    err = initFabricState.Init();

    SuccessOrQuit(err, "initFabricState.Init failed\n");

    initiatorEng.Pw = (const uint8_t *)mInitPW;

    initiatorEng.PwLen = (uint16_t)strlen(mInitPW);

onReconfig:

    responderEng.Init();

    err = respFabricState.Init();

    setAllowedResponderConfigs(responderEng);

    SuccessOrQuit(err, "respFabricState.Init failed\n");

    respFabricState.PairingCode = mRespPW;

    // =========== Start PASE InitiatorStep1 ==============================

    msgBuf = PacketBuffer::New();

    VerifyOrQuit(msgBuf != NULL, "PacketBuffer::New() failed");

    // Initiator generates and sends PASE Initiator Step 1 message.

    err = initiatorEng.GenerateInitiatorStep1(msgBuf, ProposedConfig(), initNodeId, respNodeId, sessionKeyId, encType, pwSrc, &initFabricState, mConfirmKey);

    if (IsExpectedError("Initiator:GenerateInitiatorStep1", err))

        goto onExpectedError;

    SuccessOrQuit(err, "WeavePASEEngine::GenerateInitiatorStep1 failed\n");

    // =========== Initiator Sends InitiatorStep1 to Responder ============

    mMutator->MutateMessage(INITIATOR_STEP_1, msgBuf);

    if (LogMessageData())

    {

        printf("Initiator->Responder: InitiatorStep1 Message (%d bytes)\n", msgBuf->DataLength());

        DumpMemory(msgBuf->Start(), msgBuf->DataLength(), "  ", 16);

    }

    // =========== Responder Processes PASE InitiatorStep1 ================

    err = responderEng.ProcessInitiatorStep1(msgBuf, respNodeId, initNodeId, &respFabricState);

    if (IsExpectedError(INITIATOR_STEP_1, err))

        goto onExpectedError;

    if (ExpectReconfig())

    {

        VerifyOrQuit(err == WEAVE_ERROR_PASE_RECONFIGURE_REQUIRED, "WEAVE_ERROR_PASE_RECONFIG_REQUIRED error expected");

        PacketBuffer::Free(msgBuf);

        msgBuf = NULL;

        // =========== Responder generates PASE ResponderReconfigMessage ==

        {

            msgBuf = PacketBuffer::New();

            err = responderEng.GenerateResponderReconfigure(msgBuf);

            SuccessOrQuit(err, "WeavePASEEngine::GenerateResponderReconfigure failed\n");

            // Reset PASE Engines

            responderEng.Reset();

        }

            // ========== Responder sends ResponderReconfig Message ============

        mMutator->MutateMessage(RESPONDER_RECONFIGURE, msgBuf);

        if (LogMessageData())

        {

            printf("Responder->Initiator: ResponderReconfig Message (%d bytes)\n", msgBuf->DataLength());

            DumpMemory(msgBuf->Start(), msgBuf->DataLength(), "  ", 16);

        }

        // =========== Initiator processes PASE ResponderReconfig =========

        {

            uint32_t tempProposedConfig = mProposedConfig;

            err = initiatorEng.ProcessResponderReconfigure(msgBuf, mProposedConfig);

            if (IsExpectedError("Initiator:ProcessResponderReconfigure", err))

            {

                mProposedConfig = tempProposedConfig;

                goto onExpectedError;

            }

            SuccessOrQuit(err, "WeavePASEEngine::ProcessResponderReconfigure failed\n");

            PacketBuffer::Free(msgBuf);

            msgBuf = NULL;

        }

        respFabricState.Shutdown();

        mExpectReconfig = false;

        goto onReconfig;

    } else {

        VerifyOrQuit(err != WEAVE_ERROR_PASE_RECONFIGURE_REQUIRED, "Unexpected reconfig!");

    }

    SuccessOrQuit(err, "WeavePASEEngine::ProcessInitiatorStep1 failed\n");

    PacketBuffer::Free(msgBuf);

    msgBuf = NULL;

    // =========== Responder Generates ResponderStep1 and ResponderStep2 ==

    {

        msgBuf = PacketBuffer::New();

        err = responderEng.GenerateResponderStep1(msgBuf);

        SuccessOrQuit(err, "WeavePASEEngine::GenerateResponderStep1 failed\n");

        // Responder generates and sends PASE Responder Step 2 message.

        msgBuf2 = PacketBuffer::New();

        err = responderEng.GenerateResponderStep2(msgBuf2);

        SuccessOrQuit(err, "WeavePASEEngine::GenerateResponderStep2 failed\n");

    }

    // =========== Responder sends ResponderStep1 ==========================

    mMutator->MutateMessage(RESPONDER_STEP_1, msgBuf);

    if (LogMessageData())

    {

        printf("Responder->Initiator: ResponderStep1 Message (%d bytes)\n", msgBuf->DataLength());

        DumpMemory(msgBuf->Start(), msgBuf->DataLength(), "  ", 16);

    }

    // =========== Responder sends ResponderStep2 ==========================

    mMutator->MutateMessage(RESPONDER_STEP_2, msgBuf2);

    if (LogMessageData())

    {

        printf("Responder->Initiator: ResponderStep2 Message (%d bytes)\n", msgBuf2->DataLength());

        DumpMemory(msgBuf2->Start(), msgBuf2->DataLength(), "  ", 16);

    }

    // =========== Initator Parses ResponderStep1 and ResponderStep2 ======

    {

        // Initiator receives and processes PASE Responder Step 1 message.

        err = initiatorEng.ProcessResponderStep1(msgBuf);

        if (IsExpectedError(RESPONDER_STEP_1, err))

            goto onExpectedError;

        SuccessOrQuit(err, "WeavePASEEngine::ProcessResponderStep1 failed\n");

        PacketBuffer::Free(msgBuf);

        msgBuf = NULL;

        // Initiator receives and processes PASE Responder Step 2 message.

        err = initiatorEng.ProcessResponderStep2(msgBuf2);

        if (IsExpectedError(RESPONDER_STEP_2, err))

            goto onExpectedError;

        SuccessOrQuit(err, "WeavePASEEngine::ProcessResponderStep2 failed\n");

        PacketBuffer::Free(msgBuf2);

        msgBuf2 = NULL;

    }

    // =========== Initator Generates InitatorStep2 ===========================

    {

        msgBuf = PacketBuffer::New();

        err = initiatorEng.GenerateInitiatorStep2(msgBuf);

        SuccessOrQuit(err, "WeavePASEEngine::GenerateInitiatorStep2 failed\n");

    }

    //=========== Initator Sends InitatorStep2 ============================

    mMutator->MutateMessage(INITIATOR_STEP_2, msgBuf);

    if (LogMessageData())

    {

        printf("Initiator->Responder: InitatorStep2 Message (%d bytes)\n", msgBuf->DataLength());

        DumpMemory(msgBuf->Start(), msgBuf->DataLength(), "  ", 16);

    }

    // =========== Responder Parses InitatorStep2 =========================

    {

        err = responderEng.ProcessInitiatorStep2(msgBuf);

        PacketBuffer::Free(msgBuf);

        msgBuf = NULL;

        if (IsExpectedError(INITIATOR_STEP_2, err))

            goto onExpectedError;

        if (expectSuccess)

            SuccessOrQuit(err, "WeavePASEEngine::ProcessInitiatorStep2 failed\n");

        else if (mConfirmKey)

        {

            VerifyOrQuit(err == WEAVE_ERROR_KEY_CONFIRMATION_FAILED, "Expected error from WeavePASEEngine::ProcessInitiatorStep2\n");

            return;

        }

    }

    if (mConfirmKey)

    {

        // ========== Responder Forms ResponderKeyConfirm =================

        {

            msgBuf = PacketBuffer::New();

            err = responderEng.GenerateResponderKeyConfirm(msgBuf);

            SuccessOrQuit(err, "WeavePASEEngine::GenerateResponderKeyConfirm failed\n");

        }

        // ========== Responder Sends ResponderKeyConfirm to Responder ====

        mMutator->MutateMessage(RESPONDER_KEY_CONFIRM, msgBuf);

        if (LogMessageData())

        {

            printf("Responder->Initiator: ResponderKeyConfirm Message (%d bytes)\n", msgBuf->DataLength());

            DumpMemory(msgBuf->Start(), msgBuf->DataLength(), "  ", 16);

        }

        // ========== Initiator Processes ResponderKeyConfirm =============

        {

            err = initiatorEng.ProcessResponderKeyConfirm(msgBuf);

            if (IsExpectedError(RESPONDER_KEY_CONFIRM, err))

                goto onExpectedError;

            SuccessOrQuit(err, "WeavePASEEngine::ProcessResponderKeyConfirm failed\n");

            PacketBuffer::Free(msgBuf);

            msgBuf = NULL;

        }

    }

    VerifyOrQuit(initiatorEng.State == WeavePASEEngine::kState_InitiatorDone, "Initiator state != Done\n");

    VerifyOrQuit(responderEng.State == WeavePASEEngine::kState_ResponderDone, "Responder state != Done\n");

    VerifyOrQuit(initiatorEng.SessionKeyId == responderEng.SessionKeyId, "Initiator SessionKeyId != Responder SessionKeyId\n");

    VerifyOrQuit(initiatorEng.EncryptionType == responderEng.EncryptionType, "Initiator EncryptionType != Responder EncryptionType\n");

    VerifyOrQuit(initiatorEng.PerformKeyConfirmation == responderEng.PerformKeyConfirmation, "Initiator SessionKeyId != Responder SessionKeyId\n");

    err = initiatorEng.GetSessionKey(initiatorKey);

    SuccessOrQuit(err, "WeavePASEEngine::GetSessionKey() failed\n");

    err = responderEng.GetSessionKey(responderKey);

    SuccessOrQuit(err, "WeavePASEEngine::GetSessionKey() failed\n");

    VerifyOrQuit(memcmp(initiatorKey->AES128CTRSHA1.DataKey, responderKey->AES128CTRSHA1.DataKey, WeaveEncryptionKey_AES128CTRSHA1::DataKeySize) == 0,

                 "Data key mismatch\n");

    VerifyOrQuit(memcmp(initiatorKey->AES128CTRSHA1.IntegrityKey, responderKey->AES128CTRSHA1.IntegrityKey, WeaveEncryptionKey_AES128CTRSHA1::IntegrityKeySize) == 0,

                 "Integrity key mismatch\n");

    // Shutdown the Initiator/Responder FabricState objects

    err = initFabricState.Shutdown();

    SuccessOrQuit(err, "initFabricState.Shutdown failed\n");

    err = respFabricState.Shutdown();

    SuccessOrQuit(err, "respFabricState.Shutdown failed\n");

onExpectedError:

    PacketBuffer::Free(msgBuf);

    msgBuf = NULL;

    PacketBuffer::Free(msgBuf2);

    msgBuf2 = NULL;

    initiatorEng.Shutdown();

    responderEng.Shutdown();

    initFabricState.Shutdown();

    respFabricState.Shutdown();

    if (LogMessageData())

        printf("Test Complete: %s\n", TestName());

}