# Michael Cohen <scudette@users.sourceforge.net>
# David Collett <daveco@users.sourceforge.net>
#
# ******************************************************
#  Version: FLAG $Version: 0.87-pre1 Date: Thu Jun 12 00:48:38 EST 2008$
# ******************************************************
#
# * This program is free software; you can redistribute it and/or
# * modify it under the terms of the GNU General Public License
# * as published by the Free Software Foundation; either version 2
# * of the License, or (at your option) any later version.
# *
# * This program is distributed in the hope that it will be useful,
# * but WITHOUT ANY WARRANTY; without even the implied warranty of
# * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# * GNU General Public License for more details.
# *
# * You should have received a copy of the GNU General Public License
# * along with this program; if not, write to the Free Software
# * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
# ******************************************************
""" This module provides support for compressed file formats such as Zip and Gzip.

The scanner recurses into zip files, executing the scanner factory train on files within the ZIP file.

Note that the scanner reads the central directory to recover
compressed file offsets with in the zip file. The carver does not and
finds zip file contents where ever they appear with
in the zip -
hopefull the both return the same thing, but in the case of a multi
file zip file
 the carver will work better than the scanner.
"""
import os.path,sys,posixpath
import pyflag.pyflaglog as pyflaglog
from pyflag.Scanner import *
import zipfile,gzip,tarfile, zlib
from pyflag.FileSystem import File
import pyflag.Magic as Magic
import time,re,os
import StringIO
import pyflag.Scanner as Scanner
import gzip
import pyflag.Store as Store
import FileFormats.Zip as Zip
import pyflag.Time as Time

class ZipScan(GenScanFactory):
    """ Recurse into Zip Files """
    order=99
    default = True
    group = "CompressedFile"

    def scan(self, fd, scanners, type, mime, cookie, **args):
        pass

    class Scan:
        types = (
            'application/(x-)?zip',
            )

        def external_process(self,fd):
            """ This is run on the extracted file """
            pyflaglog.log(pyflaglog.VERBOSE_DEBUG, "Decompressing Zip File %s" % fd.inode)
            cache_key = "%s:%s" % (self.case , self.fd.inode)

            ## Try to read the fd as a zip file
            z = zipfile.ZipFile(fd)

            pathname, inode, inode_id = self.ddfs.lookup(inode = self.inode)

            ## retrieve evidence timezone, this is necessary because zip files
            ## store time in localtime
            evidence_tz = Time.get_evidence_tz_name(self.case, self.fd)
            
            ## List all the files in the zip file:
            dircount = 0
            inodes = []
            namelist = z.namelist()
            for i in range(len(namelist)):
                ## Add the file into the VFS
                try:
                    ## Convert the time to case timezone
                    t = Time.convert(z.infolist()[i].date_time, case=self.case, evidence_tz=evidence_tz)
                except:
                    t=0

                ## If the entry corresponds to just a directory we ignore it.
                if not posixpath.basename(namelist[i]): continue

                info = z.infolist()[i]
                inode = "%s|Z%s:%s" % (self.inode,info.header_offset, info.compress_size)
                inodes.append(inode)
                
                inode_id = self.ddfs.VFSCreate(None,
                                               inode,DB.expand("%s/%s",(pathname,namelist[i])),
                                               size=info.file_size,
                                               mtime=t, _fast=True)
                
            for inode in inodes:
                ## Now call the scanners on this new file (FIXME limit
                ## the recursion level here)
                fd = self.ddfs.open(inode_id = inode_id)
                Scanner.scanfile(self.ddfs,fd,self.factories)

class GZScan(GenScanFactory):
    """ Decompress Gzip files """
    group = "CompressedFile"
    
    class Drawer(Scanner.Drawer):
        description = "Compressed file support"
        name = "Compressed File"
        group = "CompressedFile"
        default = False

    def find_gzipped_filename(self, fd, type):
        match = re.search(type,'was "([^"]+)"')
        if match:
            return match.groups(1)

        if fd.urn.endswith(".gz"):
            return original_filename[:-3]

        return "Uncompressed"

    ## For gziped files, we convert them to Image stream to provide
    ## seekable access - This does duplicate data but otherwise we
    ## would need to use a temporary file anyway
    def scan(self, fd, scanners, type, mime, cookie, **args):
        if "gzip" in type:
            new_path = "%s/%s" % (fd.urn, self.find_gzipped_filename(fd, type))
            new_fd = CacheManager.AFF4_MANAGER.create_cache_fd(self.case, new_path)
            gz = gzip.GzipFile(fileobj=fd, mode='r')
            while 1:
                data = gz.read(1024*1024)
                if not data:  break
                new_fd.write(data)

            new_fd.close()
            ## Now scan the new fd
            Scanner.scan_inode(self.case, new_fd.inode_id,
                               scanners)

    class Scan:
        """ If we hit a gzip file, we just create a new Inode entry in the VFS """
        types = (
            'application/x-gzip' ,
            )
        
        def __init__(self, inode,ddfs,outer,factories=None,fd=None):
            ScanIfType.__init__(self, inode,ddfs,outer,factories,fd=fd)
            self.filename = None

        def process(self, data, metadata=None):
            ScanIfType.process(self,data,metadata)
            if not self.boring_status and not self.filename:
                ## We need to find the name of the original uncompressed
                ## file so we can set a sensible VFS file name. This is
                ## the algorithm used:
                ## 1) We try to decompress the first data block from the file to see if the original name is in the header
                ## 2) Failing this we check if the inodes filename ends with .gz
                ## 3) Failing that, we call the new file "data"
                m = Magic.MagicResolver()
                magic, type_mime = m.find_inode_magic(self.case, inode_id=self.fd.inode_id,
                                                      data=data[:1024])

        def finish(self):
            if self.filename:
                self.ddfs.VFSCreate(self.inode,"G0",self.filename)

                new_inode="%s|G0" % (self.inode)
                ## Scan the new file using the scanner train:
                fd=self.ddfs.open(inode=new_inode)
                Scanner.scanfile(self.ddfs,fd,self.factories)

class TarScan(GenScanFactory):
    """ Recurse into Tar Files """
    order=99
    default = True
    depends = [ 'TypeScan' ]
    group = 'CompressedFile'

    def destroy(self):
        pass
    
    class Scan:
        types = (
            'application/x-tar',
            )

        def external_process(self,fd):
            """ This is run on the extracted file """
            #Get a TarFile object - We must access a complete file
            #here
            fd.cache()
            tar=tarfile.TarFile(fileobj=fd)
            
            ## List all the files in the tar file:
            inodes = []
            dircount = 0
            namelist = tar.getnames()
            for i in range(len(namelist)):

                ## If the entry corresponds to just a directory we ignore it.
                if not os.path.basename(namelist[i]): continue
                
                ## Add the file into the VFS
                self.ddfs.VFSCreate(
                    self.inode,"T%s" % i,namelist[i],
                    size=tar.getmember(namelist[i]).size,
                    _mtime=tar.getmember(namelist[i]).mtime,
          