/*

 * The Sleuth Kit

 *

 * Brian Carrier [carrier <at> sleuthkit [dot] org]

 * Copyright (c) 2019-2020 Brian Carrier.  All Rights reserved

 * Copyright (c) 2018-2019 BlackBag Technologies.  All Rights reserved

 *

 * This software is distributed under the Common Public License 1.0

 */

#include "tsk_apfs.hpp"

#include "tsk/libtsk.h"

#include "tsk/fs/tsk_apfs.hpp"

#include "tsk/fs/tsk_fs_i.h"

#include "tsk/img/legacy_cache.h"

#include <stdexcept>

APFSPool::APFSPool(std::vector<img_t>&& imgs, apfs_block_num nx_block_num)

    : TSKPool(std::forward<std::vector<img_t>>(imgs)),

      _nx_block_num{nx_block_num} {

  if (_members.size() != 1) {

    throw std::runtime_error(

        "Only single physical store APFS pools are currently supported");

  }

  // If we're scanning for the latest NXSB then we need to start with the

  // last known good NXSB first

  if (_nx_block_num == APFS_POOL_NX_BLOCK_LATEST) {

    _nx_block_num = APFS_POOL_NX_BLOCK_LAST_KNOWN_GOOD;

  }

  std::tie(_img, _offset) = _members[0];

  auto nxsb = nx(true);

  // Update the base members

  _uuid = nxsb->uuid();

  _block_size = nxsb->block_size();

  _dev_block_size = _img->sector_size;

  _num_blocks = nxsb->num_blocks();

  // Check to see if we need to scan for a newer pool

  if (nx_block_num == APFS_POOL_NX_BLOCK_LATEST) {

    const auto versions = known_versions();

    if (versions.empty()) {

      _nx_block_num = APFS_POOL_NX_BLOCK_LAST_KNOWN_GOOD;

      if (tsk_verbose) {

        tsk_fprintf(stderr,

                    "APFSPool: No checkpoint superblocks found.  Attempting to "

                    "fall back to last known good superblock\n");

      }

    } else {

      const auto highest = std::max_element(

          versions.begin(), versions.end(),

          [](const auto& a, const auto& b) { return a.xid < b.xid; });

      // No need to do anything, we're already the highest version

      if (highest->xid != nxsb->xid()) {

        _nx_block_num = highest->nx_block_num;

        try {

          nxsb = nx(true);

        } catch (const std::runtime_error&) {

          // Fallback to last known good block if the latest block is not valid

          _nx_block_num = APFS_POOL_NX_BLOCK_LAST_KNOWN_GOOD;

          nxsb = nx(true);

        }

      }

    }

  }

  _vol_blocks = nxsb->volume_blocks();

  _num_vols = static_cast<int>(_vol_blocks.size());

  // If the software crypto bit is not set, then either hardware crypto is used

  // or there are no volumes that are encrypted.

  if (bit_is_set(nxsb->sb()->flags, APFS_NXSB_FLAGS_CRYPTO_SW) == false) {

    // We need to check each volume to determine if any of them have encryption

    // enabled.

    for (const auto& volume : volumes()) {

      if (volume.encrypted()) {

        _hw_crypto = true;

        break;

      }

    }

  }

}

std::unique_ptr<APFSSuperblock> APFSPool::nx(bool validate) const {

  auto nxsb = std::make_unique<APFSSuperblock>((*this), _nx_block_num);

  if (validate && nxsb->validate_checksum() == false) {

    throw std::runtime_error("NXSB object checksum failed");

  }

  return nxsb;

}

std::vector<APFSFileSystem> APFSPool::volumes() const {

  std::vector<APFSFileSystem> v{};

  v.reserve(_vol_blocks.size());

  for (const auto block : _vol_blocks) {

    v.emplace_back((*this), block);

  }

  return v;

}

ssize_t APFSPool::read(uint64_t address, char* buf, size_t buf_size) const

    noexcept {

  return tsk_img_read(_img, address + _offset, buf, buf_size);

}

const std::vector<APFSPool::nx_version> APFSPool::known_versions() const {

  std::vector<nx_version> vers{};

  const auto nxsb = nx();

  const auto sb = nxsb->sb();

  for (auto addr = sb->chkpt_desc_base_addr;

       addr < sb->chkpt_desc_base_addr + sb->chkpt_desc_block_count; addr++) {

    APFSObject obj{(*this), addr};

    if (obj.obj_type() != APFS_OBJ_TYPE_SUPERBLOCK ||

        obj.oid() != nxsb->oid()) {

      continue;

    }

    if (!obj.validate_checksum()) {

      continue;

    }

    vers.emplace_back(nx_version{addr, obj.xid()});

  }

  return vers;

}

const std::vector<APFSPool::range> APFSPool::unallocated_ranges() const {

  return nx()->unallocated_ranges();

}

void APFSPool::clear_cache() noexcept {

  _block_cache.clear();

  auto cache = static_cast<LegacyCache*>(reinterpret_cast<IMG_INFO*>(_img)->cache);

  cache->lock();

  cache->clear();

  cache->unlock();

}