From ec8db1694bd84024b2cc773095601ed895835588 Mon Sep 17 00:00:00 2001 From: Jack Andersen Date: Mon, 29 Jun 2015 14:07:46 -1000 Subject: [PATCH] FST building/traversal in place --- include/DiscBase.hpp | 91 ++++++++++++++---- include/aes.hpp | 96 +------------------ lib/DiscBase.cpp | 37 +++++++- lib/DiscGCN.cpp | 42 ++++++++- lib/DiscWii.cpp | 22 ++++- lib/NODLib.cpp | 1 + lib/aes.cpp | 213 ++++++++++++++++++++++++++++++------------- main.cpp | 16 ++++ 8 files changed, 336 insertions(+), 182 deletions(-) diff --git a/include/DiscBase.hpp b/include/DiscBase.hpp index 03b7f33..044f931 100644 --- a/include/DiscBase.hpp +++ b/include/DiscBase.hpp @@ -4,6 +4,7 @@ #include #include #include +#include #include "Util.hpp" #include "IDiscIO.hpp" #include "IFileIO.hpp" @@ -35,22 +36,9 @@ public: } }; - struct FSTNode - { - uint32_t typeAndNameOff; - uint32_t offset; - uint32_t length; - FSTNode(IDiscIO::IReadStream& s) - { - s.read(this, 12); - typeAndNameOff = SBig(typeAndNameOff); - offset = SBig(offset); - length = SBig(length); - } - }; - struct IPartReadStream { + virtual void seek(size_t offset, int whence=SEEK_SET)=0; virtual size_t read(void* buf, size_t length)=0; }; @@ -59,6 +47,18 @@ public: virtual size_t write(void* buf, size_t length)=0; }; + class FSTNode + { + uint32_t typeAndNameOffset; + uint32_t offset; + uint32_t length; + public: + inline bool isDir() const {return SBig(typeAndNameOffset) >> 24;} + inline uint32_t getNameOffset() const {return SBig(typeAndNameOffset) & 0xffffff;} + inline uint32_t getOffset() const {return SBig(offset);} + inline uint32_t getLength() const {return SBig(length);} + }; + class IPartition { public: @@ -77,18 +77,27 @@ public: NODE_DIRECTORY }; private: + friend class IPartition; const IPartition& m_parent; Kind m_kind; - std::string m_name; std::unique_ptr m_hddFile; size_t m_discOffset; size_t m_discLength; + std::string m_name; + + std::vector::const_iterator m_childrenBegin; + std::vector::const_iterator m_childrenEnd; public: - Node(const IPartition& parent, bool isDir, const char* name) - : m_parent(parent), m_kind(isDir ? NODE_DIRECTORY : NODE_FILE), m_name(name) {} + Node(const IPartition& parent, const FSTNode& node, const char* name) + : m_parent(parent), + m_kind(node.isDir() ? NODE_DIRECTORY : NODE_FILE), + m_discOffset(node.getOffset()), + m_discLength(node.getLength()), + m_name(name) {} inline Kind getKind() const {return m_kind;} + inline const std::string& getName() const {return m_name;} std::unique_ptr beginReadStream() const { if (m_kind != NODE_FILE) @@ -98,10 +107,37 @@ public: } return m_parent.beginReadStream(m_discOffset); } + inline std::vector::const_iterator rawBegin() const {return m_childrenBegin;} + inline std::vector::const_iterator rawEnd() const {return m_childrenEnd;} + + class DirectoryIterator : std::iterator + { + friend class Node; + std::vector::const_iterator m_it; + DirectoryIterator(std::vector::const_iterator&& it) + : m_it(std::move(it)) {} + public: + inline bool operator!=(const DirectoryIterator& other) {return m_it != other.m_it;} + inline DirectoryIterator& operator++() + { + if (m_it->m_kind == NODE_DIRECTORY) + m_it = m_it->rawEnd(); + else + ++m_it; + return *this; + } + inline const Node& operator*() {return *m_it;} + }; + inline DirectoryIterator begin() const {return DirectoryIterator(rawBegin());} + inline DirectoryIterator end() const {return DirectoryIterator(rawEnd());} }; protected: - std::vector m_files; - void parseFST(); + uint32_t m_dolOff; + uint32_t m_fstOff; + uint32_t m_fstSz; + uint32_t m_apploaderOff; + std::vector m_nodes; + void parseFST(IPartReadStream& s); const DiscBase& m_parent; Kind m_kind; @@ -111,6 +147,7 @@ public: : m_parent(parent), m_kind(kind), m_offset(offset) {} inline Kind getKind() const {return m_kind;} virtual std::unique_ptr beginReadStream(size_t offset=0) const=0; + inline const Node& getFSTRoot() const {return m_nodes[0];} }; protected: @@ -120,8 +157,22 @@ protected: public: DiscBase(std::unique_ptr&& dio); virtual bool commit()=0; - inline Header getHeader() const {return m_header;} + inline const Header& getHeader() const {return m_header;} inline const IDiscIO& getDiscIO() const {return *m_discIO.get();} + inline const IPartition* getDataPartition() const + { + for (const std::unique_ptr& part : m_partitions) + if (part->getKind() == IPartition::PART_DATA) + return part.get(); + return nullptr; + } + inline const IPartition* getUpdatePartition() const + { + for (const std::unique_ptr& part : m_partitions) + if (part->getKind() == IPartition::PART_UPDATE) + return part.get(); + return nullptr; + } }; } diff --git a/include/aes.hpp b/include/aes.hpp index 80531fb..962d580 100644 --- a/include/aes.hpp +++ b/include/aes.hpp @@ -1,11 +1,10 @@ -#ifndef __AES_HPP_ -#define __AES_HPP_ +#ifndef __AES_HPP__ +#define __AES_HPP__ #include #include #include #include -#include namespace NOD { @@ -18,97 +17,8 @@ public: virtual void setKey(const uint8_t* key)=0; }; -class SoftwareAES : public IAES -{ -protected: - uint8_t fbsub[256]; - uint8_t rbsub[256]; - uint8_t ptab[256], ltab[256]; - uint32_t ftable[256]; - uint32_t rtable[256]; - uint32_t rco[30]; - - /* Parameter-dependent data */ - - int Nk, Nb, Nr; - uint8_t fi[24], ri[24]; - uint32_t fkey[120]; - uint32_t rkey[120]; - - - uint8_t bmul(uint8_t x, uint8_t y); - uint32_t SubByte(uint32_t a); - uint8_t product(uint32_t x, uint32_t y); - uint32_t InvMixCol(uint32_t x); - uint8_t ByteSub(uint8_t x); - void gentables(void); - void gkey(int nb, int nk, const uint8_t* key); - void _encrypt(uint8_t* buff); - void _decrypt(uint8_t* buff); - -public: - void encrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, size_t len); - void decrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, size_t len); - void setKey(const uint8_t* key); -}; - -#if __AES__ - -#include - -class NiAES : public IAES -{ - __m128i m_ekey[11]; - __m128i m_dkey[11]; -public: - void encrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, size_t len) - { - __m128i feedback,data; - size_t i,j; - if (len%16) - len = len/16+1; - else - len /= 16; - feedback = _mm_loadu_si128((__m128i*)iv); - for (i=0 ; i NewAES(); } -#endif //__AES_HPP_ +#endif //__AES_HPP__ diff --git a/lib/DiscBase.cpp b/lib/DiscBase.cpp index ea47801..645b1df 100644 --- a/lib/DiscBase.cpp +++ b/lib/DiscBase.cpp @@ -8,12 +8,39 @@ DiscBase::DiscBase(std::unique_ptr&& dio) { } -void DiscBase::IPartition::parseFST() +void DiscBase::IPartition::parseFST(IPartReadStream& s) { - char buf[1024]; - std::unique_ptr s = beginReadStream(); - s->read(buf, 1024); - printf(""); + std::unique_ptr fst(new uint8_t[m_fstSz]); + s.seek(m_fstOff); + s.read(fst.get(), m_fstSz); + + const FSTNode* nodes = (FSTNode*)fst.get(); + + /* Root node indicates the count of all contained nodes */ + uint32_t nodeCount = nodes[0].getLength(); + const char* names = (char*)fst.get() + 12 * nodeCount; + m_nodes.clear(); + m_nodes.reserve(nodeCount); + + /* Construct nodes */ + for (uint32_t n=0 ; n"); + } + + /* Setup dir-child iterators */ + for (std::vector::iterator it=m_nodes.begin(); + it != m_nodes.end(); + ++it) + { + Node& node = *it; + if (node.m_kind == Node::NODE_DIRECTORY) + { + node.m_childrenBegin = it + 1; + node.m_childrenEnd = m_nodes.begin() + node.m_discLength; + } + } } } diff --git a/lib/DiscGCN.cpp b/lib/DiscGCN.cpp index d569890..d483df5 100644 --- a/lib/DiscGCN.cpp +++ b/lib/DiscGCN.cpp @@ -3,10 +3,50 @@ namespace NOD { +class PartitionGCN : public DiscBase::IPartition +{ +public: + PartitionGCN(const DiscGCN& parent, Kind kind, size_t offset) + : IPartition(parent, kind, offset) + { + /* GCN-specific header reads */ + std::unique_ptr s = beginReadStream(0x420); + uint32_t vals[3]; + s->read(vals, 12); + m_dolOff = SBig(vals[0]); + m_fstOff = SBig(vals[1]); + m_fstSz = SBig(vals[2]); + + /* Yay files!! */ + parseFST(*s.get()); + } + + class PartReadStream : public DiscBase::IPartReadStream + { + const PartitionGCN& m_parent; + std::unique_ptr m_dio; + + public: + PartReadStream(const PartitionGCN& parent, size_t offset) + : m_parent(parent) + {m_dio = m_parent.m_parent.getDiscIO().beginReadStream(offset);} + void seek(size_t offset, int whence) + {m_dio->seek(offset, whence);} + size_t read(void* buf, size_t length) + {return m_dio->read(buf, length);} + }; + + std::unique_ptr beginReadStream(size_t offset) const + { + return std::unique_ptr(new PartReadStream(*this, offset)); + } +}; + DiscGCN::DiscGCN(std::unique_ptr&& dio) : DiscBase(std::move(dio)) { - + /* One lone partition for GCN */ + m_partitions.emplace_back(new PartitionGCN(*this, IPartition::PART_DATA, 0)); } bool DiscGCN::commit() diff --git a/lib/DiscWii.cpp b/lib/DiscWii.cpp index 8547756..5dc71dd 100644 --- a/lib/DiscWii.cpp +++ b/lib/DiscWii.cpp @@ -213,7 +213,16 @@ public: aes->setKey(COMMON_KEY); aes->decrypt(iv, m_ticket.encKey, m_decKey, 16); - parseFST(); + /* Wii-specific header reads (now using title key to decrypt) */ + std::unique_ptr ds = beginReadStream(0x420); + uint32_t vals[3]; + ds->read(vals, 12); + m_dolOff = SBig(vals[0]) << 2; + m_fstOff = SBig(vals[1]) << 2; + m_fstSz = SBig(vals[2]) << 2; + + /* Yay files!! */ + parseFST(*ds.get()); } class PartReadStream : public DiscBase::IPartReadStream @@ -240,6 +249,17 @@ public: size_t block = m_offset / 0x7c00; m_dio = m_parent.m_parent.getDiscIO().beginReadStream(m_baseOffset + block * 0x8000); } + void seek(size_t offset, int whence) + { + if (whence == SEEK_SET) + m_offset = offset; + else if (whence == SEEK_CUR) + m_offset += offset; + else + return; + size_t block = m_offset / 0x7c00; + m_dio->seek(m_baseOffset + block * 0x8000); + } size_t read(void* buf, size_t length) { size_t cacheOffset = m_offset % 0x7c00; diff --git a/lib/NODLib.cpp b/lib/NODLib.cpp index bb1c34a..53dd96f 100644 --- a/lib/NODLib.cpp +++ b/lib/NODLib.cpp @@ -52,6 +52,7 @@ std::unique_ptr OpenDiscFromImage(const char* path, bool& isWii) if (!discIO) { + fclose(fp); throw std::runtime_error("'" + std::string(path) + "' is not a valid image"); return std::unique_ptr(); } diff --git a/lib/aes.cpp b/lib/aes.cpp index 0c6c0db..e905321 100644 --- a/lib/aes.cpp +++ b/lib/aes.cpp @@ -12,6 +12,7 @@ #include //#include #include +#include namespace NOD { @@ -30,13 +31,13 @@ namespace NOD static const uint8_t InCo[4] = {0xB, 0xD, 0x9, 0xE}; /* Inverse Coefficients */ -static uint32_t pack(const uint8_t* b) +static inline uint32_t pack(const uint8_t* b) { /* pack bytes into a 32-bit Word */ return ((uint32_t)b[3] << 24) | ((uint32_t)b[2] << 16) | ((uint32_t)b[1] << 8) | (uint32_t)b[0]; } -static void unpack(uint32_t a, uint8_t* b) +static inline void unpack(uint32_t a, uint8_t* b) { /* unpack bytes from a word */ b[0] = (uint8_t)a; @@ -45,7 +46,7 @@ static void unpack(uint32_t a, uint8_t* b) b[3] = (uint8_t)(a >> 24); } -static uint8_t xtime(uint8_t a) +static inline uint8_t xtime(uint8_t a) { uint8_t b; @@ -57,6 +58,40 @@ static uint8_t xtime(uint8_t a) return a; } +class SoftwareAES : public IAES +{ +protected: + uint8_t fbsub[256]; + uint8_t rbsub[256]; + uint8_t ptab[256], ltab[256]; + uint32_t ftable[256]; + uint32_t rtable[256]; + uint32_t rco[30]; + + /* Parameter-dependent data */ + + int Nk, Nb, Nr; + uint8_t fi[24], ri[24]; + uint32_t fkey[120]; + uint32_t rkey[120]; + + + uint8_t bmul(uint8_t x, uint8_t y); + uint32_t SubByte(uint32_t a); + uint8_t product(uint32_t x, uint32_t y); + uint32_t InvMixCol(uint32_t x); + uint8_t ByteSub(uint8_t x); + void gentables(void); + void gkey(int nb, int nk, const uint8_t* key); + void _encrypt(uint8_t* buff); + void _decrypt(uint8_t* buff); + +public: + void encrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, size_t len); + void decrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, size_t len); + void setKey(const uint8_t* key); +}; + uint8_t SoftwareAES::bmul(uint8_t x, uint8_t y) { /* x.y= AntiLog(Log(x) + Log(y)) */ @@ -454,76 +489,127 @@ void SoftwareAES::encrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, si #if __AES__ -static inline __m128i AES_128_ASSIST (__m128i temp1, __m128i temp2) +#include + +class NiAES : public IAES { - __m128i temp3; - temp2 = _mm_shuffle_epi32 (temp2 ,0xff); - temp3 = _mm_slli_si128 (temp1, 0x4); - temp1 = _mm_xor_si128 (temp1, temp3); - temp3 = _mm_slli_si128 (temp3, 0x4); - temp1 = _mm_xor_si128 (temp1, temp3); - temp3 = _mm_slli_si128 (temp3, 0x4); - temp1 = _mm_xor_si128 (temp1, temp3); - temp1 = _mm_xor_si128 (temp1, temp2); - return temp1; -} + __m128i m_ekey[11]; + __m128i m_dkey[11]; +public: + void encrypt(uint8_t* iv, const uint8_t* inbuf, uint8_t* outbuf, size_t len) + { + __m128i feedback,data; + size_t i,j; + if (len%16) + len = len/16+1; + else + len /= 16; + feedback = _mm_loadu_si128((__m128i*)iv); + for (i=0 ; i NewAES() { +#if __AES__ if (HAS_AES_NI == -1) { unsigned int a,b,c,d; @@ -537,6 +623,9 @@ std::unique_ptr NewAES() return std::unique_ptr(new NiAES); else return std::unique_ptr(new SoftwareAES); +#else + return std::unique_ptr(new SoftwareAES); +#endif } } diff --git a/main.cpp b/main.cpp index 0d85ec0..b1abde8 100644 --- a/main.cpp +++ b/main.cpp @@ -14,6 +14,22 @@ int main(int argc, char* argv[]) if (!disc) return -1; + const NOD::DiscBase::IPartition* dataPart = disc->getDataPartition(); + if (dataPart) + { + for (const NOD::DiscBase::IPartition::Node& node : dataPart->getFSTRoot()) + { + if (node.getKind() == NOD::DiscBase::IPartition::Node::NODE_FILE) + printf("FILE: %s\n", node.getName().c_str()); + else if (node.getKind() == NOD::DiscBase::IPartition::Node::NODE_DIRECTORY) + { + printf("DIR: %s\n", node.getName().c_str()); + for (const NOD::DiscBase::IPartition::Node& subnode : node) + printf("SUBFILE: %s\n", subnode.getName().c_str()); + } + } + } + return 0; }