path: root/tools/archiver/main.cpp

#include <library/cpp/archive/yarchive.h>
#include <library/cpp/deprecated/mapped_file/mapped_file.h>
#include <library/cpp/digest/md5/md5.h>
#include <library/cpp/getopt/small/last_getopt.h>

#include <util/folder/dirut.h>
#include <util/folder/filelist.h>
#include <util/folder/path.h>
#include <util/generic/vector.h>
#include <util/generic/yexception.h>
#include <util/memory/blob.h>
#include <util/stream/file.h>
#include <util/string/cast.h>
#include <util/string/escape.h>
#include <util/string/hex.h>
#include <util/string/subst.h>
#include <util/system/filemap.h> 

#include <cstring>

namespace {
    class TStringArrayOutput: public IOutputStream {
    public:
        TStringArrayOutput(IOutputStream* slave, size_t stride)
            : Slave(*slave)
            , Stride(stride)
        {
            Buf.reserve(stride);
        }
        void DoFinish() override {
            WriteBuf();
            Flush();
        }
        void DoWrite(const void* data, size_t len) override {
            for (const char* p = (const char*)data; len > 0; ++p, --len) {
                Buf.append(*p);
                if (Buf.size() == Stride)
                    WriteBuf();
            }
        }

    private:
        void WriteBuf() {
            Slave << '"' << Buf << "\",\n"sv;
            Buf.clear();
        }

    private:
        IOutputStream& Slave;
        const size_t Stride;
        TString Buf;
    };

    class THexOutput: public IOutputStream {
    public:
        inline THexOutput(IOutputStream* slave)
            : Slave_(slave) 
        {
        }

        ~THexOutput() override {
        }

        inline IOutputStream* Slave() const noexcept {
            return Slave_; 
        }

    private:
        void DoFinish() override {
            Slave_->Write('\n'); 
            Slave_->Flush(); 
        }

        void DoWrite(const void* data, size_t len) override {
            const char* b = (const char*)data;

            while (len) {
                const unsigned char c = *b;
                char buf[12];
                char* tmp = buf;

                if (Count_ % Columns == 0) { 
                    *tmp++ = ' ';
                    *tmp++ = ' ';
                    *tmp++ = ' ';
                    *tmp++ = ' ';
                }

                if (Count_ && Count_ % Columns != 0) { 
                    *tmp++ = ',';
                    *tmp++ = ' ';
                }

                *tmp++ = '0';
                *tmp++ = 'x';
                tmp = HexEncode(&c, 1, tmp);

                if ((Count_ % Columns) == (Columns - 1)) { 
                    *tmp++ = ',';
                    *tmp++ = '\n';
                }

                Slave_->Write(buf, tmp - buf); 

                --len;
                ++b;
                ++Count_; 
            }
        }

    private:
        // width in source chars
        static const size_t Columns = 10; 
        ui64 Count_ = 0; 
        IOutputStream* Slave_ = nullptr;
    };

    struct TYasmOutput: public IOutputStream {
        inline TYasmOutput(IOutputStream* out, const TString& base)
            : Out_(out) 
            , Base_(base) 
        {
            *Out_ << "global " << Base_ << "\n"; 
            *Out_ << "global " << Base_ << "Size\n\nSECTION .rodata\n\n"; 
            *Out_ << Base_ << ":\n"; 
        }

        ~TYasmOutput() override {
        }

        void DoFinish() override {
            *Out_ << Base_ << "Size:\ndd " << Count_ << '\n'; 

            *Out_ << "%ifidn __OUTPUT_FORMAT__,elf64\n";
            *Out_ << "size " << Base_ << " " << Count_ << "\n";
            *Out_ << "size " << Base_ << "Size 4\n";
            *Out_ << "%endif\n";
        }

        void DoWrite(const void* data, size_t len) override {
            Count_ += len; 

            const unsigned char* p = (const unsigned char*)data;

            while (len) {
                const size_t step = Min<size_t>(len, 100);

                *Out_ << "db " << (int)*p++; 

                for (size_t i = 1; i < step; ++i) {
                    *Out_ << ',' << (int)*p++; 
                }

                *Out_ << '\n'; 

                len -= step;
            }
        }

        IOutputStream* Out_ = nullptr;
        const TString Base_;
        ui64 Count_ = 0; 
    };

    struct TCOutput: public THexOutput {
        inline TCOutput(IOutputStream* out, const TString& base)
            : THexOutput(out)
            , B(base)
        {
            *Slave() << "static_assert(sizeof(unsigned int) == 4, \"ups, unsupported platform\");\n\nextern \"C\" {\nextern const unsigned char " << B << "[] = {\n";
        }

        ~TCOutput() override {
        }

        void DoFinish() override {
            *Slave() << "\n};\nextern const unsigned int " << B << "Size = sizeof(" << B << ") / sizeof(" << B << "[0]);\n}\n";
        }

        const TString B;
    };

    struct TCStringOutput: public IOutputStream {
        inline TCStringOutput(IOutputStream* out, const TString& base)
            : O(out)
            , B(base)
        {
            *O << "static_assert(sizeof(unsigned int) == 4, \"ups, unsupported platform\");\n\nextern \"C\" {\nextern const unsigned char " << B << "[] = \n";
        }

        ~TCStringOutput() override {
        }

        void DoWrite(const void* data, size_t len) override {
            *O << TString((const char*)data, len).Quote() << '\n';
        }

        void DoFinish() override {
            //*O << ";\nextern const unsigned char* " << B << " = (const unsigned char*)" << B << "Array;\n";
            *O << ";\nextern const unsigned int " << B << "Size = sizeof(" << B << ") / sizeof(" << B << "[0]) - 1;\n}\n";
        }

        IOutputStream* O = nullptr;
        const TString B;
    };

    struct TMyFileComparator {
        bool operator()(const TString& fname1, const TString& fname2) const {
            if (fname1 == fname2) {
                return false;
            }
            if (const auto* savedResultPtr = SavedResults.FindPtr(std::make_pair(fname1, fname2))) {
                return *savedResultPtr < 0;
            }
            TMemoryMap mmap1(fname1, TMemoryMap::oRdOnly);
            TMemoryMap mmap2(fname2, TMemoryMap::oRdOnly);
            mmap1.SetSequential();
            mmap2.SetSequential();
            Y_ASSERT(mmap1.Length() == mmap2.Length());
            TMemoryMap::TMapResult mapResult1 = mmap1.Map(0, mmap1.Length());
            TMemoryMap::TMapResult mapResult2 = mmap2.Map(0, mmap2.Length());
            Y_ASSERT(mapResult1.MappedSize() == mapResult2.MappedSize());
            int res = memcmp(mapResult1.MappedData(), mapResult2.MappedData(), mapResult1.MappedSize());
            mmap1.Unmap(mapResult1);
            mmap2.Unmap(mapResult2);
            SavedResults[std::make_pair(fname1, fname2)] = res;
            SavedResults[std::make_pair(fname2, fname1)] = -res;
            return res < 0;
        }

        mutable THashMap<std::pair<TString, TString>, int> SavedResults;
    };

    struct TDuplicatesMap {
        void Add(const TString& fname, const TString& rname) {
            Y_ENSURE(!InitialFillingDone);
            FileNames.push_back(fname);
            FileNameToRecordName[fname] = rname;
        }

        void Finish() {
            Y_ENSURE(!InitialFillingDone);
            InitialFillingDone = true;
            TMap<i64, TVector<TString>> bySize;
            for (const TString& fname: FileNames) {
                TFile file(fname, OpenExisting | RdOnly);
                bySize[file.GetLength()].push_back(fname);
            }
            for (const auto& bySizeElement: bySize) {
                if (bySizeElement.second.size() > 1) {
                    TMap<TString, TVector<TString>, TMyFileComparator> byContents;
                    for (const TString& fname: bySizeElement.second) {
                        byContents[fname].push_back(fname);
                    }
                    for (const auto& byContentsElement: byContents) {
                        if (byContentsElement.second.size() > 1) {
                            const TString& rootName = byContentsElement.second.front();
                            const TString& rootRecordName = FileNameToRecordName[rootName];
                            for (const TString& fname: byContentsElement.second) {
                                if (fname != rootName) {
                                    Synonyms[FileNameToRecordName[fname]] = rootRecordName;
                                }
                            }
                        }
                    }
                }
            }
            FileNames.clear();
            FileNameToRecordName.clear();
        }

        bool InitialFillingDone = false;
        TVector<TString> FileNames;
        THashMap<TString, TString> FileNameToRecordName;
        THashMap<TString, TString> Synonyms;
    };

    struct TDeduplicationArchiveWriter {
        TDeduplicationArchiveWriter(const TDuplicatesMap& duplicatesMap, IOutputStream* out, bool compress)
            : DuplicatesMap(duplicatesMap)
            , Writer(out, compress)
        {}

        void Finish() {
            Writer.Finish();
        }

        const TDuplicatesMap& DuplicatesMap;
        TArchiveWriter Writer;
    };
}

static inline TAutoPtr<IOutputStream> OpenOutput(const TString& url) {
    if (url.empty()) {
        return new TBuffered<TUnbufferedFileOutput>(8192, Duplicate(1));
    } else {
        return new TBuffered<TUnbufferedFileOutput>(8192, url);
    }
}

static inline bool IsDelim(char ch) noexcept {
    return ch == '/' || ch == '\\';
}

static inline TString GetFile(const TString& s) {
    const char* e = s.end();
    const char* b = s.begin();
    const char* c = e - 1;

    while (c != b && !IsDelim(*c)) {
        --c;
    }

    if (c != e && IsDelim(*c)) {
        ++c;
    }

    return TString(c, e - c);
}

static inline TString Fix(TString f) {
    if (!f.empty() && IsDelim(f[f.size() - 1])) {
        f.pop_back();
    }

    return f;
}

static bool Quiet = false;

static inline void Append(IOutputStream& w, const TString& fname, const TString& rname) {
    TMappedFileInput in(fname);

    if (!Quiet) { 
        Cerr << "--> " << rname << Endl;
    } 

    TransferData((IInputStream*)&in, &w);
}

static inline void Append(TDuplicatesMap& w, const TString& fname, const TString& rname) {
    w.Add(fname, rname);
}

static inline void Append(TDeduplicationArchiveWriter& w, const TString& fname, const TString& rname) {
    if (!Quiet) { 
        Cerr << "--> " << rname << Endl;
    } 

    if (const TString* rootRecordName = w.DuplicatesMap.Synonyms.FindPtr(rname)) {
        w.Writer.AddSynonym(*rootRecordName, rname);
    } else {
        TMappedFileInput in(fname);
        w.Writer.Add(rname, &in);
    }
}

namespace {
    struct TRec {
        bool Recursive = false; 
        TString Key;
        TString Path;
        TString Prefix;

        TRec() = default; 

        inline void Fix() {
            ::Fix(Path);
            ::Fix(Prefix);
        }

        template <typename T>
        inline void Recurse(T& w) const {
            if (IsDir(Path)) {
                DoRecurse(w, "/");
            } else {
                Append(w, Path, Key.size() ? Key : Prefix + "/" + GetFile(Path));
            }
        }

        template <typename T>
        inline void DoRecurse(T& w, const TString& off) const {
            {
                TFileList fl;

                const char* name;
                const TString p = Path + off;

                fl.Fill(p, true);

                while ((name = fl.Next())) {
                    const TString fname = p + name;
                    const TString rname = Prefix + off + name;

                    Append(w, fname, rname);
                }
            }

            if (Recursive) {
                TDirsList dl;

                const char* name;
                const TString p = Path + off;

                dl.Fill(p, true);

                while ((name = dl.Next())) {
                    if (strcmp(name, ".") && strcmp(name, "..")) {
                        DoRecurse(w, off + name + "/");
                    }
                }
            }
        }
    };
}

static TString CutFirstSlash(const TString& fileName) {
    if (fileName[0] == '/') {
        return fileName.substr(1);
    } else {
        return fileName;
    }
}

struct TMappingReader { 
    TMemoryMap Map; 
    TBlob Blob; 
    TArchiveReader Reader; 
 
    TMappingReader(const TString& archive)
        : Map(archive) 
        , Blob(TBlob::FromMemoryMapSingleThreaded(Map, 0, Map.Length())) 
        , Reader(Blob) 
    { 
    } 
}; 
 
static void UnpackArchive(const TString& archive, const TFsPath& dir = TFsPath()) {
    TMappingReader mappingReader(archive); 
    const TArchiveReader& reader = mappingReader.Reader; 
    const size_t count = reader.Count();
    for (size_t i = 0; i < count; ++i) {
        const TString key = reader.KeyByIndex(i);
        const TString fileName = CutFirstSlash(key);
        if (!Quiet) { 
            Cerr << archive << " --> " << fileName << Endl;
        } 
        const TFsPath path(dir / fileName);
        path.Parent().MkDirs();
        TAutoPtr<IInputStream> in = reader.ObjectByKey(key);
        TFixedBufferFileOutput out(path);
        TransferData(in.Get(), &out);
        out.Finish();
    }
}

static void ListArchive(const TString& archive, bool cutSlash) {
    TMappingReader mappingReader(archive); 
    const TArchiveReader& reader = mappingReader.Reader; 
    const size_t count = reader.Count();
    for (size_t i = 0; i < count; ++i) {
        const TString key = reader.KeyByIndex(i);
        TString fileName = key;
        if (cutSlash) {
            fileName = CutFirstSlash(key);
        }
        Cout << fileName << Endl;
    }
}

static void ListArchiveMd5(const TString& archive, bool cutSlash) {
    TMappingReader mappingReader(archive); 
    const TArchiveReader& reader = mappingReader.Reader; 
    const size_t count = reader.Count();
    for (size_t i = 0; i < count; ++i) {
        const TString key = reader.KeyByIndex(i);
        TString fileName = key;
        if (cutSlash) {
            fileName = CutFirstSlash(key);
        }
        char md5buf[33];
        Cout << fileName << '\t' << MD5::Stream(reader.ObjectByKey(key).Get(), md5buf) << Endl;
    }
}

int main(int argc, char** argv) {
    NLastGetopt::TOpts opts; 
    opts.AddHelpOption('?'); 
    opts.SetTitle( 
        "Archiver\n" 
        "Docs: https://wiki.yandex-team.ru/Development/Poisk/arcadia/tools/archiver"
    ); 

    bool hexdump = false; 
    opts.AddLongOption('x', "hexdump", "Produce hexdump") 
        .NoArgument() 
        .Optional() 
        .StoreValue(&hexdump, true); 
 
    size_t stride = 0; 
    opts.AddLongOption('s', "segments", "Produce segmented C strings array of given size") 
        .RequiredArgument("<size>") 
        .Optional() 
        .DefaultValue("0") 
        .StoreResult(&stride); 
 
    bool cat = false; 
    opts.AddLongOption('c', "cat", "Do not store keys (file names), just cat uncompressed files") 
        .NoArgument() 
        .Optional() 
        .StoreValue(&cat, true); 
 
    bool doNotZip = false; 
    opts.AddLongOption('p', "plain", "Do not use compression") 
        .NoArgument() 
        .Optional() 
        .StoreValue(&doNotZip, true); 
 
    bool deduplicate = false;
    opts.AddLongOption("deduplicate", "Turn on file-wise deduplication")
        .NoArgument()
        .Optional()
        .StoreValue(&deduplicate, true);

    bool unpack = false; 
    opts.AddLongOption('u', "unpack", "Unpack archive into current directory") 
        .NoArgument() 
        .Optional() 
        .StoreValue(&unpack, true); 
 
    bool list = false; 
    opts.AddLongOption('l', "list", "List files in archive") 
        .NoArgument() 
        .Optional() 
        .StoreValue(&list, true); 
 
    bool cutSlash = true;
    opts.AddLongOption("as-is", "somewhy slash is cutted by default in list; with this option key will be shown as-is")
        .NoArgument()
        .Optional()
        .StoreValue(&cutSlash, false);

    bool listMd5 = false; 
    opts.AddLongOption('m', "md5", "List files in archive with MD5 sums") 
        .NoArgument() 
        .Optional() 
        .StoreValue(&listMd5, true); 
 
    bool recursive = false; 
    opts.AddLongOption('r', "recursive", "Read all files under each directory, recursively") 
        .NoArgument() 
        .Optional() 
        .StoreValue(&recursive, true); 
 
    Quiet = false; 
    opts.AddLongOption('q', "quiet", "Do not output progress to stderr") 
        .NoArgument() 
        .Optional() 
        .StoreValue(&Quiet, true); 
 
    TString prepend;
    opts.AddLongOption('z', "prepend", "Prepend string to output") 
        .RequiredArgument("<prefix>") 
        .StoreResult(&prepend); 
 
    TString append;
    opts.AddLongOption('a', "append", "Append string to output") 
        .RequiredArgument("<suffix>") 
        .StoreResult(&append); 
 
    TString outputf;
    opts.AddLongOption('o', "output", "Output to file instead stdout") 
        .RequiredArgument("<file>") 
        .StoreResult(&outputf); 
 
    TString unpackDir;
    opts.AddLongOption('d', "unpackdir", "Unpack destination directory") 
        .RequiredArgument("<dir>") 
        .DefaultValue(".") 
        .StoreResult(&unpackDir); 
 
    TString yasmBase;
    opts.AddLongOption('A', "yasm", "Output dump is yasm format") 
        .RequiredArgument("<base>") 
        .StoreResult(&yasmBase); 
 
    TString cppBase;
    opts.AddLongOption('C', "cpp", "Output dump is C/C++ format") 
        .RequiredArgument("<base>") 
        .StoreResult(&cppBase); 
 
    TString forceKeys;
    opts.AddLongOption('k', "keys", "Set explicit list of keys for elements")
        .RequiredArgument("<keys>")
        .StoreResult(&forceKeys);

    opts.SetFreeArgDefaultTitle("<file>"); 
    opts.SetFreeArgsMin(1); 
    NLastGetopt::TOptsParseResult optsRes(&opts, argc, argv); 
 
    SubstGlobal(append, "\\n", "\n");
    SubstGlobal(prepend, "\\n", "\n");

    TVector<TRec> recs;
    const auto& files = optsRes.GetFreeArgs(); 

    TVector<TStringBuf> keys;
    if (forceKeys.size())
        StringSplitter(forceKeys).Split(':').SkipEmpty().Collect(&keys);

    if (keys.size() && keys.size() != files.size()) {
        Cerr << "Invalid number of keys=" << keys.size() << " (!= number of files=" << files.size() << ")" << Endl;
        return 1;
    }

    for (size_t i = 0; i < files.size(); ++i) {
        const auto& path = files[i];
        size_t off = 0;
#ifdef _win_
        if (path[0] > 0 && isalpha(path[0]) && path[1] == ':')
            off = 2; // skip drive letter ("d:")
#endif               // _win_
        const size_t pos = path.find(':', off);
        TRec cur;
        cur.Path = path.substr(0, pos);
        if (pos != TString::npos)
            cur.Prefix = path.substr(pos + 1);
        if (keys.size())
            cur.Key = keys[i];
        cur.Recursive = recursive;
        cur.Fix();
        recs.push_back(cur);
    }

    try {
        if (listMd5) {
            for (const auto& rec: recs) {
                ListArchiveMd5(rec.Path, cutSlash);
            }
        } else if (list) {
            for (const auto& rec: recs) {
                ListArchive(rec.Path, cutSlash);
            }
        } else if (unpack) {
            const TFsPath dir(unpackDir);
            for (const auto& rec: recs) {
                UnpackArchive(rec.Path, dir);
            }
        } else {
            TAutoPtr<IOutputStream> outf(OpenOutput(outputf));
            IOutputStream* out = outf.Get();
            THolder<IOutputStream> hexout;

            if (hexdump) {
                hexout.Reset(new THexOutput(out));
                out = hexout.Get();
            } else if (stride) {
                hexout.Reset(new TStringArrayOutput(out, stride));
                out = hexout.Get();
            } else if (yasmBase) {
                hexout.Reset(new TYasmOutput(out, yasmBase));
                out = hexout.Get();
            } else if (cppBase) {
                hexout.Reset(new TCStringOutput(out, cppBase));
                out = hexout.Get();
            }

            outf->Write(prepend.data(), prepend.size());

            if (cat) {
                for (const auto& rec: recs) {
                    rec.Recurse(*out);
                }
            } else {
                TDuplicatesMap duplicatesMap;
                if (deduplicate) {
                    for (const auto& rec: recs) {
                        rec.Recurse(duplicatesMap);
                    }
                }
                duplicatesMap.Finish();
                TDeduplicationArchiveWriter w(duplicatesMap, out, !doNotZip);
                for (const auto& rec: recs) {
                    rec.Recurse(w);
                }
                w.Finish();
            }

            try {
                out->Finish();
            } catch (...) {
            }

            outf->Write(append.data(), append.size());
        }
    } catch (...) {
        Cerr << CurrentExceptionMessage() << Endl;
        return 1;
    }

    return 0;
}