/*

 * 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

 */

#pragma once

#if HAVE_CONFIG_H

#include "tsk/tsk_config.h"

#endif

#include "tsk/util/crypto.hpp"

#include "apfs_fs.h"

#include "tsk_apfs.hpp"

class APFSJObject {

  using jit = APFSJObjBtreeNode::iterator;

  using child_entry = struct {

    std::string name;

    apfs_dir_record rec;

  };

  using extent = struct {

    uint64_t offset;

    uint64_t phys;

    uint64_t len;

    uint64_t crypto_id;

  };

  using inline_xattr = struct {

    std::string name;

    std::string data;

  };

  using nonres_xattr = struct {

    std::string name;

    uint64_t oid;

    uint64_t size;

    uint64_t allocated_size;

    uint64_t crypto_id;

  };

  apfs_inode _inode{};

  std::vector<child_entry> _children{};

  std::vector<extent> _extents{};

  std::vector<inline_xattr> _inline_xattrs{};

  std::vector<nonres_xattr> _nonres_xattrs{};

  std::string _name{};

  uint64_t _size{};

  uint64_t _size_on_disk{};

  bool _is_clone{};

  void add_entry(const jit::value_type &);

 public:

  using key_type = APFSJObjKey;

  APFSJObject() = default;

  APFSJObject(const std::pair<jit, jit> &);

  APFSJObject(const jit &, const jit &);

  APFSJObject(const APFSJObject &) = default;

  APFSJObject &operator=(const APFSJObject &) = default;

  APFSJObject(APFSJObject &&) = default;

  APFSJObject &operator=(APFSJObject &&) = default;

  inline bool valid() const noexcept {

    return _inode.private_id != 0 || !_extents.empty();

  }

  inline auto child_count() const noexcept { return _children.size(); }

  inline const apfs_inode &inode() const noexcept { return _inode; }

  inline const std::string &name() const noexcept { return _name; }

  inline const std::vector<extent> &extents() const noexcept {

    return _extents;

  }

  inline const std::vector<inline_xattr> &inline_xattrs() const noexcept {

    return _inline_xattrs;

  }

  inline const std::vector<nonres_xattr> &nonres_xattrs() const noexcept {

    return _nonres_xattrs;

  }

  inline const std::vector<child_entry> &children() const noexcept {

    return _children;

  }

  inline uint64_t size() const noexcept { return _size; }

  inline uint64_t size_on_disk() const noexcept { return _size_on_disk; }

  inline bool is_clone() const noexcept { return _is_clone; }

};

class APFSJObjTree {

  using jit = APFSJObjBtreeNode::iterator;

 protected:

  struct crypto {

#ifdef HAVE_LIBCRYPTO

    std::unique_ptr<aes_xts_decryptor> decryptor{};

#endif

    std::unique_ptr<uint8_t[]> key{};

    std::string password{};

    crypto(const APFSFileSystem::crypto_info_t &crypto);

  } _crypto;

  APFSObjectBtreeNode _obj_root;

  APFSJObjBtreeNode _jobj_root;

  uint64_t _root_tree_oid;

  inline auto jobjs(uint64_t oid) const {

    return _jobj_root.find_range(

        oid, [](const auto &key, const auto &b) noexcept->int64_t {

          const auto akey = key.template as<APFSJObject::key_type>();

          return akey->oid() - b;

        });

  }

  APFSJObjTree(const APFSFileSystem &vol);

 public:

  APFSJObjTree(const APFSPool &pool, apfs_block_num obj_omap,

               uint64_t root_tree_oid,

               const APFSFileSystem::crypto_info_t &crypto);

  APFSJObjTree(APFSJObjTree &&) = default;

  inline APFSJObject obj(uint64_t oid) const { return {jobjs(oid)}; }

  void set_snapshot(uint64_t snap_xid);

  // iterator stuff

  class iterator;

  iterator begin() const;

  iterator end() const;

};

class APFSJObjTree::iterator {

 public:

  using iterator_category = std::forward_iterator_tag;

  using difference_type = uint32_t;

  using value_type = APFSJObject;

  using reference = const value_type &;

  using pointer = const value_type *;

 protected:

  const APFSJObjTree *_tree;

  jit _next;

  APFSJObject _jobj;

  iterator(const APFSJObjTree *tree) noexcept;

  iterator(const APFSJObjTree *tree, uint64_t oid);

 public:

  iterator() = default;

  iterator(const iterator &) = default;

  iterator &operator=(const iterator &) = default;

  iterator(iterator &&) = default;

  iterator &operator=(iterator &&) = default;

  inline reference operator*() const noexcept { return _jobj; }

  inline pointer operator->() const noexcept { return &_jobj; }

  iterator &operator++();

  inline value_type operator++(int) {

    value_type copy{_jobj};

    this->operator++();

    return copy;

  }

  inline bool operator==(const iterator &rhs) const noexcept {

    return (_tree == rhs._tree && _next == rhs._next);

  }

  inline bool operator!=(const iterator &rhs) const noexcept {

    return !this->operator==(rhs);

  }

  friend APFSJObjTree;

};