/*
 * Copyright 2019, The Android Open Source Project
 *
 * 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.
 */

#define LOG_TAG "WritableIdentityCredential"

#include "WritableIdentityCredential.h"

#include <android/hardware/identity/support/IdentityCredentialSupport.h>

#include <android-base/logging.h>
#include <android-base/stringprintf.h>

#include <cppbor.h>
#include <cppbor_parse.h>

#include <utility>

#include "IdentityCredentialStore.h"

#include "FakeSecureHardwareProxy.h"

namespace aidl::android::hardware::identity {

using ::android::base::StringPrintf;
using ::std::optional;
using namespace ::android::hardware::identity;

bool WritableIdentityCredential::initialize() {
    if (!hwProxy_->initialize(testCredential_)) {
        LOG(ERROR) << "hwProxy->initialize() failed";
        return false;
    }
    startPersonalizationCalled_ = false;
    firstEntry_ = true;

    return true;
}

// Used when updating a credential. Returns false on failure.
bool WritableIdentityCredential::initializeForUpdate(
        const vector<uint8_t>& encryptedCredentialKeys) {
    if (!hwProxy_->initializeForUpdate(testCredential_, docType_, encryptedCredentialKeys)) {
        LOG(ERROR) << "hwProxy->initializeForUpdate() failed";
        return false;
    }
    startPersonalizationCalled_ = false;
    firstEntry_ = true;

    return true;
}

WritableIdentityCredential::~WritableIdentityCredential() {}

ndk::ScopedAStatus WritableIdentityCredential::getAttestationCertificate(
        const vector<uint8_t>& attestationApplicationId,
        const vector<uint8_t>& attestationChallenge, vector<Certificate>* outCertificateChain) {
    if (getAttestationCertificateAlreadyCalled_) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_FAILED,
                "Error attestation certificate previously generated"));
    }
    getAttestationCertificateAlreadyCalled_ = true;

    if (attestationChallenge.empty()) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_INVALID_DATA, "Challenge can not be empty"));
    }

    optional<vector<uint8_t>> certChain;
    if (attestationKeyBlob_ && attestationCertificateChain_) {
        certChain = hwProxy_->createCredentialKeyUsingRkp(
                attestationChallenge, attestationApplicationId, *attestationKeyBlob_,
                attestationCertificateChain_->at(0));
    } else {
        certChain = hwProxy_->createCredentialKey(attestationChallenge, attestationApplicationId);
    }

    if (!certChain) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_FAILED,
                "Error generating attestation certificate chain"));
    }

    optional<vector<vector<uint8_t>>> certs = support::certificateChainSplit(certChain.value());
    if (!certs) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_FAILED,
                "Error splitting chain into separate certificates"));
    }

    *outCertificateChain = vector<Certificate>();
    for (vector<uint8_t>& cert : certs.value()) {
        Certificate c;
        c.encodedCertificate = std::move(cert);
        outCertificateChain->push_back(std::move(c));
    }

    for (const vector<uint8_t>& cert : *attestationCertificateChain_) {
        Certificate c;
        c.encodedCertificate = cert;
        outCertificateChain->push_back(std::move(c));
    }

    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus WritableIdentityCredential::setExpectedProofOfProvisioningSize(
        int32_t expectedProofOfProvisioningSize) {
    expectedProofOfProvisioningSize_ = expectedProofOfProvisioningSize;
    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus WritableIdentityCredential::startPersonalization(
        int32_t accessControlProfileCount, const vector<int32_t>& entryCounts) {
    if (startPersonalizationCalled_) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_FAILED, "startPersonalization called already"));
    }
    startPersonalizationCalled_ = true;

    numAccessControlProfileRemaining_ = accessControlProfileCount;
    remainingEntryCounts_ = entryCounts;
    entryNameSpace_ = "";

    signedDataAccessControlProfiles_ = cppbor::Array();
    signedDataNamespaces_ = cppbor::Map();
    signedDataCurrentNamespace_ = cppbor::Array();

    if (!hwProxy_->startPersonalization(accessControlProfileCount, entryCounts, docType_,
                                        expectedProofOfProvisioningSize_)) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_FAILED, "eicStartPersonalization"));
    }

    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus WritableIdentityCredential::addAccessControlProfile(
        int32_t id, const Certificate& readerCertificate, bool userAuthenticationRequired,
        int64_t timeoutMillis, int64_t secureUserId,
        SecureAccessControlProfile* outSecureAccessControlProfile) {
    if (numAccessControlProfileRemaining_ == 0) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_INVALID_DATA,
                "numAccessControlProfileRemaining_ is 0 and expected non-zero"));
    }

    if (accessControlProfileIds_.find(id) != accessControlProfileIds_.end()) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_INVALID_DATA,
                "Access Control Profile id must be unique"));
    }
    accessControlProfileIds_.insert(id);

    if (id < 0 || id >= 32) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_INVALID_DATA,
                "Access Control Profile id must be non-negative and less than 32"));
    }

    // Spec requires if |userAuthenticationRequired| is false, then |timeoutMillis| must also
    // be zero.
    if (!userAuthenticationRequired && timeoutMillis != 0) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_INVALID_DATA,
                "userAuthenticationRequired is false but timeout is non-zero"));
    }

    optional<vector<uint8_t>> mac = hwProxy_->addAccessControlProfile(
            id, readerCertificate.encodedCertificate, userAuthenticationRequired, timeoutMillis,
            secureUserId);
    if (!mac) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_FAILED, "eicAddAccessControlProfile"));
    }

    SecureAccessControlProfile profile;
    profile.id = id;
    profile.readerCertificate = readerCertificate;
    profile.userAuthenticationRequired = userAuthenticationRequired;
    profile.timeoutMillis = timeoutMillis;
    profile.secureUserId = secureUserId;
    profile.mac = mac.value();
    cppbor::Map profileMap;
    profileMap.add("id", profile.id);
    if (profile.readerCertificate.encodedCertificate.size() > 0) {
        profileMap.add("readerCertificate",
                       cppbor::Bstr(profile.readerCertificate.encodedCertificate));
    }
    if (profile.userAuthenticationRequired) {
        profileMap.add("userAuthenticationRequired", profile.userAuthenticationRequired);
        profileMap.add("timeoutMillis", profile.timeoutMillis);
    }
    signedDataAccessControlProfiles_.add(std::move(profileMap));

    numAccessControlProfileRemaining_--;

    *outSecureAccessControlProfile = profile;
    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus WritableIdentityCredential::beginAddEntry(
        const vector<int32_t>& accessControlProfileIds, const string& nameSpace, const string& name,
        int32_t entrySize) {
    if (numAccessControlProfileRemaining_ != 0) {
        LOG(ERROR) << "numAccessControlProfileRemaining_ is " << numAccessControlProfileRemaining_
                   << " and expected zero";
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_INVALID_DATA,
                "numAccessControlProfileRemaining_ is not zero"));
    }

    // Ensure passed-in profile ids reference valid access control profiles
    for (const int32_t id : accessControlProfileIds) {
        if (accessControlProfileIds_.find(id) == accessControlProfileIds_.end()) {
            return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                    IIdentityCredentialStore::STATUS_INVALID_DATA,
                    "An id in accessControlProfileIds references non-existing ACP"));
        }
    }

    if (remainingEntryCounts_.size() == 0) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_INVALID_DATA, "No more namespaces to add to"));
    }

    // Handle initial beginEntry() call.
    if (firstEntry_) {
        firstEntry_ = false;
        entryNameSpace_ = nameSpace;
        allNameSpaces_.insert(nameSpace);
    }

    // If the namespace changed...
    if (nameSpace != entryNameSpace_) {
        if (allNameSpaces_.find(nameSpace) != allNameSpaces_.end()) {
            return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                    IIdentityCredentialStore::STATUS_INVALID_DATA,
                    "Name space cannot be added in interleaving fashion"));
        }

        // Then check that all entries in the previous namespace have been added..
        if (remainingEntryCounts_[0] != 0) {
            return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                    IIdentityCredentialStore::STATUS_INVALID_DATA,
                    "New namespace but a non-zero number of entries remain to be added"));
        }
        remainingEntryCounts_.erase(remainingEntryCounts_.begin());
        remainingEntryCounts_[0] -= 1;
        allNameSpaces_.insert(nameSpace);

        if (signedDataCurrentNamespace_.size() > 0) {
            signedDataNamespaces_.add(entryNameSpace_, std::move(signedDataCurrentNamespace_));
            signedDataCurrentNamespace_ = cppbor::Array();
        }
    } else {
        // Same namespace...
        if (remainingEntryCounts_[0] == 0) {
            return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                    IIdentityCredentialStore::STATUS_INVALID_DATA,
                    "Same namespace but no entries remain to be added"));
        }
        remainingEntryCounts_[0] -= 1;
    }

    entryRemainingBytes_ = entrySize;
    entryNameSpace_ = nameSpace;
    entryName_ = name;
    entryAccessControlProfileIds_ = accessControlProfileIds;
    entryBytes_.resize(0);
    // LOG(INFO) << "name=" << name << " entrySize=" << entrySize;

    if (!hwProxy_->beginAddEntry(accessControlProfileIds, nameSpace, name, entrySize)) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_FAILED, "eicBeginAddEntry"));
    }

    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus WritableIdentityCredential::addEntryValue(const vector<uint8_t>& content,
                                                             vector<uint8_t>* outEncryptedContent) {
    size_t contentSize = content.size();

    if (contentSize > IdentityCredentialStore::kGcmChunkSize) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_INVALID_DATA,
                "Passed in chunk of is bigger than kGcmChunkSize"));
    }
    if (contentSize > entryRemainingBytes_) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_INVALID_DATA,
                "Passed in chunk is bigger than remaining space"));
    }

    entryBytes_.insert(entryBytes_.end(), content.begin(), content.end());
    entryRemainingBytes_ -= contentSize;
    if (entryRemainingBytes_ > 0) {
        if (contentSize != IdentityCredentialStore::kGcmChunkSize) {
            return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                    IIdentityCredentialStore::STATUS_INVALID_DATA,
                    "Retrieved non-final chunk which isn't kGcmChunkSize"));
        }
    }

    optional<vector<uint8_t>> encryptedContent = hwProxy_->addEntryValue(
            entryAccessControlProfileIds_, entryNameSpace_, entryName_, content);
    if (!encryptedContent) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_FAILED, "eicAddEntryValue"));
    }

    if (entryRemainingBytes_ == 0) {
        // TODO: ideally do do this without parsing the data (but still validate data is valid
        // CBOR).
        auto [item, _, message] = cppbor::parse(entryBytes_);
        if (item == nullptr) {
            return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                    IIdentityCredentialStore::STATUS_INVALID_DATA, "Data is not valid CBOR"));
        }
        cppbor::Map entryMap;
        entryMap.add("name", entryName_);
        entryMap.add("value", std::move(item));
        cppbor::Array profileIdArray;
        for (auto id : entryAccessControlProfileIds_) {
            profileIdArray.add(id);
        }
        entryMap.add("accessControlProfiles", std::move(profileIdArray));
        signedDataCurrentNamespace_.add(std::move(entryMap));
    }

    *outEncryptedContent = encryptedContent.value();
    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus WritableIdentityCredential::finishAddingEntries(
        vector<uint8_t>* outCredentialData, vector<uint8_t>* outProofOfProvisioningSignature) {
    if (numAccessControlProfileRemaining_ != 0) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_INVALID_DATA,
                "numAccessControlProfileRemaining_ is not 0 and expected zero"));
    }

    if (remainingEntryCounts_.size() > 1 || remainingEntryCounts_[0] != 0) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_INVALID_DATA,
                "More entry spaces remain than startPersonalization configured"));
    }

    if (signedDataCurrentNamespace_.size() > 0) {
        signedDataNamespaces_.add(entryNameSpace_, std::move(signedDataCurrentNamespace_));
    }
    cppbor::Array popArray;
    popArray.add("ProofOfProvisioning")
            .add(docType_)
            .add(std::move(signedDataAccessControlProfiles_))
            .add(std::move(signedDataNamespaces_))
            .add(testCredential_);
    vector<uint8_t> encodedCbor = popArray.encode();

    if (encodedCbor.size() != expectedProofOfProvisioningSize_) {
        LOG(ERROR) << "CBOR for proofOfProvisioning is " << encodedCbor.size() << " bytes, "
                   << "was expecting " << expectedProofOfProvisioningSize_;
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_INVALID_DATA,
                StringPrintf("Unexpected CBOR size %zd for proofOfProvisioning, was expecting %zd",
                             encodedCbor.size(), expectedProofOfProvisioningSize_)
                        .c_str()));
    }

    optional<vector<uint8_t>> signatureOfToBeSigned = hwProxy_->finishAddingEntries();
    if (!signatureOfToBeSigned) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_FAILED, "eicFinishAddingEntries"));
    }

    optional<vector<uint8_t>> signature =
            support::coseSignEcDsaWithSignature(signatureOfToBeSigned.value(),
                                                encodedCbor,  // data
                                                {});          // certificateChain
    if (!signature) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_FAILED, "Error signing data"));
    }

    optional<vector<uint8_t>> encryptedCredentialKeys = hwProxy_->finishGetCredentialData(docType_);
    if (!encryptedCredentialKeys) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_FAILED,
                "Error generating encrypted CredentialKeys"));
    }
    cppbor::Array array;
    array.add(docType_);
    array.add(testCredential_);
    array.add(encryptedCredentialKeys.value());
    vector<uint8_t> credentialData = array.encode();

    *outCredentialData = credentialData;
    *outProofOfProvisioningSignature = signature.value();
    hwProxy_->shutdown();

    return ndk::ScopedAStatus::ok();
}

ndk::ScopedAStatus WritableIdentityCredential::setRemotelyProvisionedAttestationKey(
        const vector<uint8_t>& attestationKeyBlob,
        const vector<uint8_t>& attestationCertificateChain) {
    if (!hardwareInformation_.isRemoteKeyProvisioningSupported) {
        return ndk::ScopedAStatus(AStatus_fromExceptionCodeWithMessage(
                EX_UNSUPPORTED_OPERATION, "Remote key provisioning is not supported"));
    }

    if (attestationKeyBlob.empty() || attestationCertificateChain.empty()) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_FAILED,
                "Empty data passed to setRemotlyProvisionedAttestationKey"));
    }

    if (attestationKeyBlob_.has_value()) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_FAILED, "Attestation key already set"));
    }

    optional<vector<vector<uint8_t>>> certs =
            support::certificateChainSplit(attestationCertificateChain);
    if (!certs) {
        return ndk::ScopedAStatus(AStatus_fromServiceSpecificErrorWithMessage(
                IIdentityCredentialStore::STATUS_FAILED,
                "Error splitting chain into separate certificates"));
    }

    attestationKeyBlob_ = attestationKeyBlob;
    attestationCertificateChain_ = *certs;
    return ndk::ScopedAStatus::ok();
}

}  // namespace aidl::android::hardware::identity