#!/usr/bin/python2.4

##############################################################################
#
# Copyright (c) 2007 Nexedi SARL. All Rights Reserved.
#
# This software is subject to the provisions of the Zope Public License,
# Version 2.0 (ZPL).  A copy of the ZPL should accompany this distribution.
# THIS SOFTWARE IS PROVIDED "AS IS" AND ANY AND ALL EXPRESS OR IMPLIED
# WARRANTIES ARE DISCLAIMED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF TITLE, MERCHANTABILITY, AGAINST INFRINGEMENT, AND FITNESS
# FOR A PARTICULAR PURPOSE
#
##############################################################################
# Parts of this file are borrowed from Zope 2.8.8 repozo.py script.
# Essentialy "usage", "parseargs" and parts of "restore" methods.
# So it's released under the ZPL v2.0, as is Zope 2.8.8 .

"""
Usage: %(program)s [-h|--help] [-c|--config configuration_file]

  -h
  --help
    Display this help and exit.

  -c configuration_file
  --config configuration_file
    Use given file as configuration file.
    It must be a python file.
    Required if neither -h nor --help are given.
"""

import imp
import getopt
import sys
import os
from struct import pack
import shutil
from ZODB.FileStorage import FileStorage

program = sys.argv[0]

def log(message):
  print message

def parse(status_file):
  tid_log = open(status_file)
  content = {}
  last_timestamp = None

  line = tid_log.readline()
  while line != '':
    split_line = line.split(' ', 2)
    assert len(split_line) == 3, repr(split_line)
    line_timestamp, line_type, line_dict = split_line
    line_timestamp = float(line_timestamp)
    assert line_type in ('f', 'd'), repr(line_type)
    if last_timestamp is None:
      last_timestamp = line_timestamp
    else:
      assert last_timestamp < line_timestamp, '%r < %r' % (last_timestamp, line_timestamp)
    line_dict = eval(line_dict, None)
    assert isinstance(line_dict, dict), type(line_dict)
    assert len(line_dict), repr(line_dict)
    if line_type == 'd':
      for key, value in line_dict.iteritems():
        if key in content:
          assert content[key] < value, '%r < %r' % (content[key], value)
        content[key] = value
    elif line_type == 'f':
      for key, value in content.iteritems():
        assert key in line_dict, repr(key)
        assert value <= line_dict[key], '%r <= %r' % (value, line_dict[key])
      content = line_dict
    line = tid_log.readline()
  return content

READCHUNK = 10 * 1024 * 1024

def get_tid_position(filepath,last_tid):
  tid = pack('>Q', last_tid + 1)
  # Find the file position of the last completed transaction.
  fs = FileStorage(filepath, read_only=True, stop=tid)
  # Note that the FileStorage ctor calls read_index() which scans the file
  # and returns "the position just after the last valid transaction record".
  # getSize() then returns this position, which is exactly what we want,
  # because we only want to copy stuff from the beginning of the file to the
  # last valid transaction record.
  pos = fs.getSize()
  fs.close()
  return pos

def recover(data_fs_backup_path_dict, status_file):
  last_tid_dict = parse(status_file)
  for storage_id, (file_path, backup_path) in data_fs_backup_path_dict.iteritems():
    # Derived from repozo (function=do_full_backup)
    # TODO: optimise to read backup only once.
    can_restore = False
    print "Considering:", storage_id
    if os.path.exists(backup_path):
      if os.path.exists(file_path):
        print 'Both original and backup files exist for %r. If previous restoration was successful, you should delete the backup for this restoration to take place. Original: %r Backup: %r' % (storage_id, file_path, backup_path)
      else:
        print 'Only backup file is available for %r: %r. Assuming it\'s ok and restoring to %r' % (storage_id, backup_path, file_path)
        can_restore = True
    else:
      if os.path.exists(file_path):
        sys.stdout.write('Copying %r to %r... ' % (file_path, backup_path))
        sys.stdout.flush()
        fsrc = file(file_path, 'rb')
        fdst = file(backup_path, 'wb')
        # shutil.copy uses too small a buffer size, which is slow:
        shutil.copyfileobj(fsrc, fdst, READCHUNK)
        fsrc.close()
        fdst.close()
        initial_size = os.stat(file_path).st_size
        final_size = os.stat(backup_path).st_size
        if initial_size == final_size:
          can_restore = True
          print 'Done.'
        else:
          print 'Backup size %i differs from original size %i. Is the original file (%r) still in use ? Is there enough free disk space at destination (%r) ?' % (final_size, initial_size, file_path, backup_path)
      else:
        print 'Cannot find any file for %r: %r and %r do not exist.' % (storage_id, file_path, backup_path)
    if can_restore:
      tid = last_tid_dict[storage_id]
      sys.stdout.write('Scanning %r for tid %r... ' % (backup_path, tid))
      sys.stdout.flush()
      pos = get_tid_position(backup_path,last_tid_dict[storage_id])
      print 'Done. Position:', pos
      original_file = open(file_path, 'rb+')
      original_file.truncate(pos)
      original_file.close()
      print 'Truncated %s to: %s bytes (transaction %r)' % (file_path, pos, tid)
    else:
      print 'Skipping restoration of %r (%r).' % (file_path, storage_id)

def usage(code, msg=''):
  outfp = sys.stderr
  if code == 0:
    outfp = sys.stdout

  print >> outfp, __doc__ % globals()
  if msg:
    print >> outfp, msg

  sys.exit(code)

def parseargs():
  try:
    opts, args = getopt.getopt(sys.argv[1:], 'hc:',
                               ['help', 'config='])
  except getopt.error, msg:
    usage(1, msg)

  class Options:
    configuration_file_name = None
    status_file = None

  options = Options()

  for opt, arg in opts:
    if opt in ('-h', '--help'):
      usage(0)
    elif opt in ('-c', '--config'):
      options.configuration_file_name = arg

  if options.configuration_file_name is None:
    usage(1, 'Either -c or --config is required.')

  configuration_filename, ext = os.path.splitext(os.path.basename(options.configuration_file_name))
  configuration_path = os.path.dirname(options.configuration_file_name)
  if len(configuration_path):
    configuration_path = [configuration_path]
  else:
    configuration_path = sys.path
  file, path, description = imp.find_module(configuration_filename, configuration_path)
  module = imp.load_module(configuration_filename, file, path, description)
  file.close()
  try:
    options.data_fs_backup_path_dict = module.data_fs_backup_path_dict
    options.status_file = module.status_file
  except AttributeError, msg:
    usage(1, msg)
  return options

def main():
  options = parseargs()
  recover(
    data_fs_backup_path_dict=options.data_fs_backup_path_dict,
    status_file=options.status_file)

if __name__ == '__main__':
  sys.exit(main())
