// Copyright 2018 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef SRC_CONNECTIVITY_BLUETOOTH_CORE_BT_HOST_SM_UTIL_H_
#define SRC_CONNECTIVITY_BLUETOOTH_CORE_BT_HOST_SM_UTIL_H_

#include "src/connectivity/bluetooth/core/bt-host/common/byte_buffer.h"
#include "src/connectivity/bluetooth/core/bt-host/common/device_address.h"
#include "src/connectivity/bluetooth/core/bt-host/common/uint128.h"
#include "src/connectivity/bluetooth/core/bt-host/common/uint256.h"
#include "src/connectivity/bluetooth/core/bt-host/hci-spec/constants.h"
#include "src/connectivity/bluetooth/core/bt-host/sm/delegate.h"
#include "src/connectivity/bluetooth/core/bt-host/sm/error.h"
#include "src/connectivity/bluetooth/core/bt-host/sm/smp.h"
#include "src/connectivity/bluetooth/core/bt-host/sm/types.h"

namespace bt::sm::util {
// NOTE:
// All cryptographic utility functions use little-endian input/output unless explicitly noted.

// Returns the size of the SMP Packet with the given `Payload` template parameter type.
template <typename Payload>
constexpr size_t PacketSize() {
  return sizeof(Header) + sizeof(Payload);
}

// Returns a string representation of a given pairing method.
std::string PairingMethodToString(PairingMethod method);

// Returns a string representation of a PairingDelegate display method.
std::string DisplayMethodToString(Delegate::DisplayMethod method);

// Returns a string representation of a given IOCapability.
std::string IOCapabilityToString(IOCapability capability);

// Returns the HCI version of an SMP IOCapability. Returns
// hci_spec::IOCapability::kNoInputNoOutput for values not in sm::IOCapability.
hci_spec::IOCapability IOCapabilityForHci(IOCapability capability);

// Utility function to heap allocate a new PDU.
MutableByteBufferPtr NewPdu(size_t param_size);

// Returns a std::pair<initiator_value, responder_value> based on |local_value|, |peer_value| and
// the local SMP Role |role|.
template <typename T>
std::pair<T, T> MapToRoles(const T& local_value, const T& peer_value, Role role) {
  if (role == Role::kInitiator) {
    return {local_value, peer_value};
  }
  return {peer_value, local_value};
}

// Used to select the key generation method as described in Vol 3, Part H,
// 2.3.5.1 based on local and peer authentication parameters:
//   - |secure_connections|: True if Secure Connections pairing is used. False
//     means Legacy Pairing.
//   - |local_oob|: Local OOB auth data is available.
//   - |peer_oob|: Peer OOB auth data is available.
//   - |mitm_required|: True means at least one of the devices requires MITM
//     protection.
//   - |local_ioc|, |peer_ioc|: Local and peer IO capabilities.
//   - |local_initiator|: True means that the local device is the initiator and
//     |local_ioc| represents the initiator's I/O capabilities.
PairingMethod SelectPairingMethod(bool secure_connections, bool local_oob, bool peer_oob,
                                  bool mitm_required, IOCapability local_ioc, IOCapability peer_ioc,
                                  bool local_initiator);

// Implements the "Security Function 'e'" defined in Vol 3, Part H, 2.2.1.
void Encrypt(const UInt128& key, const UInt128& plaintext_data, UInt128* out_encrypted_data);

// Implements the "Confirm Value Generation" or "c1" function for LE Legacy
// Pairing described in Vol 3, Part H, 2.2.3.
//
//   |tk|: 128-bit TK value
//   |rand|: 128-bit random number
//   |preq|: 56-bit SMP "Pairing Request" PDU
//   |pres|: 56-bit SMP "Pairing Response" PDU
//   |initiator_addr|: Device address of the initiator used while establishing
//                     the connection.
//   |responder_addr|: Device address of the responder used while establishing
//                     the connection.
//
// The generated confirm value will be returned in |out_confirm_value|.
void C1(const UInt128& tk, const UInt128& rand, const ByteBuffer& preq, const ByteBuffer& pres,
        const DeviceAddress& initiator_addr, const DeviceAddress& responder_addr,
        UInt128* out_confirm_value);

// Implements the "Key Generation Function s1" to generate the STK for LE Legacy
// Pairing described in Vol 3, Part H, 2.2.4.
//
//   |tk|: 128-bit TK value
//   |r1|: 128-bit random value generated by the responder.
//   |r2|: 128-bit random value generated by the initiator.
void S1(const UInt128& tk, const UInt128& r1, const UInt128& r2, UInt128* out_stk);

// Implements the "Random Address Hash Function ah" to resolve RPAs. Described
// in Vol 3, Part H, 222.
//
//   |k|: 128-bit IRK value
//   |r|: 24-bit random part of a RPA.
//
// Returns 24 bit hash value.
uint32_t Ah(const UInt128& k, uint32_t r);

// Returns true if the given |irk| can resolve the given |rpa| using the method
// described in Vol 6, Part B, 1.3.2.3.
bool IrkCanResolveRpa(const UInt128& irk, const DeviceAddress& rpa);

// Generates a RPA using the given IRK based on the method described in Vol 6,
// Part B, 1.3.2.2.
DeviceAddress GenerateRpa(const UInt128& irk);

// Generates a static or non-resolvable private random device address.
DeviceAddress GenerateRandomAddress(bool is_static);

// Implements the AES-CMAC function defined in Vol. 3, Part H, 2.2.5.
//
//   |hash_key|: Little-endian 128-bit value used as the cipher key k in the AES-CMAC algorithm.
//   |msg|: Variable length data to be encoded by |hash_key|. This is a little-endian parameter.
//
// A return value of std::nullopt indicates the calculation failed.
std::optional<UInt128> AesCmac(const UInt128& hash_key, const ByteBuffer& msg);

// Implements the "LE Secure Connections confirm value generation function f4" per Vol. 3, Part H,
// 2.2.6.
//
//   |u|: X-coordinate of the (peer/local) ECDH public key.
//   |v|: X-coordiante of the (local/peer) ECDH public key.
//   |x|: The CMAC key.
//   |z|: 0 for pairing methods besides Passkey Entry, in which case it is one bit of the passkey.
//
// A return value of std::nullopt indicates the calculation failed.
std::optional<UInt128> F4(const UInt256& u, const UInt256& v, const UInt128& x, uint8_t z);

// Implements the "LE Secure Connections key generation function f5" per Vol. 3, Part H, 2.2.7.
//
//   |dhkey|: Diffie-Hellman key generated by both parties in Phase 1 of SC (ECDH key agreement).
//   |initiator_nonce|: Nonce value generated by the initiator to avoid replay attacks.
//   |responder_nonce|: Nonce value generated by the responder to avoid replay attacks.
//   |initiator_addr|: Device address of the pairing initiator.
//   |responder_addr|: Device address of the pairing responder.
//
// A return value of std::nullopt indicates the calculation failed. Returns an |F5Results| struct
// instead of the spec-prescribed single 256-bit value for easier client access to the MacKey/LTK.
struct F5Results {
  UInt128 mac_key;
  UInt128 ltk;
};
std::optional<F5Results> F5(const UInt256& dhkey, const UInt128& initiator_nonce,
                            const UInt128& responder_nonce, const DeviceAddress& initiator_addr,
                            const DeviceAddress& responder_addr);

// Implements the "LE Secure Connections check value generation function f6" per Vol. 3, Part H,
// 2.2.8. The semantics of each parameter depend on the pairing method and current pairing state.
// See above-noted spec section for parameter descriptions. The 24-bit IOCap parameter in the spec
// signature is separated into |auth_req|, |oob|, and |io_cap| parameters for client convenience.
//
// A return value of std::nullopt indicates the calculation failed.
std::optional<UInt128> F6(const UInt128& mackey, const UInt128& n1, const UInt128& n2,
                          const UInt128& r, AuthReqField auth_req, OOBDataFlag oob,
                          IOCapability io_cap, const DeviceAddress& a1, const DeviceAddress& a2);

// Implements the "LE Secure Connections numeric comparison value generation function g2" per Vol.
// 3, Part H, 2.2.9. The value displayed to the user should be the least significant 6
// decimal digits of the result, i.e. g2 mod 10^6.
//
//   |initiator_pubkey_x|: X-coordinate of the initiator's ECDH public key
//   |responder_pubkey_x|: X-coordinate of the responder's ECDH public key
//   |initiator_nonce|: nonce value generated by the initiator to avoid replay attacks
//   |responder_nonce|: nonce value generated by the responder to avoid replay attacks
//
// A return value of std::nullopt indicates the calculation failed.
std::optional<uint32_t> G2(const UInt256& initiator_pubkey_x, const UInt256& responder_pubkey_x,
                           const UInt128& initiator_nonce, const UInt128& responder_nonce);

// Implements the "Link key conversion function h6" per Vol. 3, Part H, 2.2.10. `w` is the
// encryption key for AES-CMAC, and `key_id` is used as the input value.
//
// A return value of std::nullopt indicates the calculation failed.
std::optional<UInt128> H6(const UInt128& w, uint32_t key_id);

// Implements the "Link key conversion function h7" per Vol. 3, Part H, 2.2.11. `salt` is the
// encryption key for AES-CMAC, and `w` is the input value.
//
// A return value of std::nullopt indicates the calculation failed.
std::optional<UInt128> H7(const UInt128& salt, const UInt128& w);

// Converts an LE LTK to a BR/EDR link key for Cross Transport Key Derivation as defined in v5.2
// Vol. 3 Part H 2.4.2.4.
//
// A return value of std::nullopt indicates the conversion failed.
std::optional<UInt128> LeLtkToBrEdrLinkKey(const UInt128& le_ltk,
                                           CrossTransportKeyAlgo hash_function);

}  // namespace bt::sm::util

#endif  // SRC_CONNECTIVITY_BLUETOOTH_CORE_BT_HOST_SM_UTIL_H_