#!/usr/bin/env python
"""
Given a .tar.gz file generated by cia-backup,
extract all messages from the embedded ring buffers.
This will result in each unique message being written
to the specified file. Each message will be NUL-terminated.

This is designed to be an easy way to generate a
corpus of real messages for testing new CIA indexing
algorithms.

  usage: extract_messages.py cia-YYYY-MM-DD.tar.gz outputfile

--Micah Dowty <micah@navi.cx>
"""

import sys; sys.path[0] += '/..'
from cia.LibCIA.Stats.Messages import MessageBuffer
import tarfile, md5

import psyco
psyco.profile()

def iterBuffers(filename):
    """Given the name of a tar file, iterates over all message buffers.
       Yields (name, buffer) tuples, where 'buffer' is already initialized
       with a file object.
       """
    archive = tarfile.open(filename)
    for item in archive:
        if item.isfile() and item.name.endswith("/_msg"):
            buffer = MessageBuffer(path="", file=archive.extractfile(item))
            buffer._init()
            yield item.name, buffer

class MessageStorage:
    numTotal = 0
    numDuplicate = 0

    def store(self, text):
        pass

class FileMessageStorage(MessageStorage):
    def __init__(self, outputfile):
        self.file = open(outputfile, "wb")
        self.memo = {}

    def store(self, text):
        self.numTotal += 1
        digest = md5.new(text).digest()
        if digest in self.memo:
            self.numDuplicate += 1
            return
        self.memo[digest] = None

        if text.find("\0") >= 0:
            raise ValueError("Message has an embedded NUL")

        self.file.write(text)
        self.file.write("\0")

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print(__doc__)
        sys.exit(1)

    archive = sys.argv[1]
    storage = FileMessageStorage(sys.argv[2])

    for bufferName, buffer in iterBuffers(archive):
        for uid, message in buffer.getLatest():
            storage.store(message.toxml().encode("utf-8"))
        print("total:%d duplicate:%d -- %s" % (
            storage.numTotal, storage.numDuplicate, bufferName))

