#!/usr/bin/python
#
# Check the free space available on a writable DVD
# Should always exit with 0 status, otherwise it indicates a serious error.
# (wrong number of arguments, Python exception...)
#
#  called:  dvd-handler <dvd-device-name> operation args
#
#  where operation is one of
#    free
#    write
#
# in case of operation ``free'' returns:
# Prints on the first output line the free space available in bytes.
# If an error occurs, prints a negative number (-errno), followed,
# on the second line, by an error message.
# 
# $Id: dvd-handler.in,v 1.7.2.3 2005/11/03 14:33:16 kerns Exp $
#

import popen2
import os
import os.path
import errno
import sys
import re
import signal
import time
import array

class disk:
# Configurable values:
   
   dvdrwmediainfo = "dvd+rw-mediainfo"
   growcmd = "growisofs"
   dvdrwformat = "dvd+rw-format"
   dd = "/bin/dd"
   margin = 10485760 # 10 mb security margin

   # Comment the following line if you want the tray to be reloaded
   # when writing ends.
   growcmd += " -use-the-force-luke=notray"

# end of configurable values

###############################################################################
#
# This class represents DVD disk informations.
# When instantiated, it needs a device name.
# Status information about the device and the disk loaded is collected only when
# asked for (for example dvd-freespace doesn't need to know the media type, and
# dvd-writepart doesn't not always need to know the free space).
#
# The following methods are implemented:
# __init__       we need that...
# __repr__       this seems to be a good idea to have.
#                Quite minimalistic implementation, though.
# __str__        For casts to string. Return the current disk information
# is_empty       Returns TRUE if the disk is empty, blank... this needs more
#                work, especially concerning non-RW media and blank vs. no
#                filesystem considerations. Here, we should also look for
#                other filesystems - probably we don't want to silently
#                overwrite UDF or ext2 or anything not mentioned in fstab...
#                (NB: I don't think it is a problem)
# free           Returns the available free space.
# write          Writes one part file to disk, either starting a new file
#                system on disk, or appending to it.
#                This method should also prepare a blank disk so that a
#                certain part of the disk is used to allow detection of a
#                used disk by all / more disk drives.
# blank          Blank the device
#
###############################################################################
   def __init__(self, devicename):
      self.device = devicename
      self.disktype = "none"
      self.diskmode = "none"
      self.diskstatus = "none"
      self.hardwaredevice = "none"
      self.pid = 0
      self.next_session = -1
      self.capacity = -1

      self.freespace_collected = 0
      self.mediumtype_collected = 0

      self.growcmd += " -quiet"

      if self.is4gbsupported():
         self.growcmd += " -use-the-force-luke=4gms"

      self.growparams = " -A 'Bacula Data' -input-charset=default -iso-level 3 -pad " + \
                        "-p 'dvd-handler / growisofs' -sysid 'BACULADATA' -R"

      return

   def __repr__(self):
      return "disk(" + self.device + ") # This is an instance of class disk"

   def __str__(self):
      if not self.freespace_collected:
         self.collect_freespace();
      if not self.mediumtype_collected:
         self.collect_mediumtype();
      
      self.me  = "Class disk, initialized with device '" + self.device + "'\n"
      self.me += "type = '" + self.disktype + "' mode='" + self.diskmode + "' status = '" + self.diskstatus + "'\n"
      self.me += " next_session = " + str(self.next_session) + " capacity = " + str(self.capacity) + "\n"
      self.me += "Hardware device is '" + self.hardwaredevice + "'\n"
      self.me += "growcmd = '" + self.growcmd + "'\ngrowparams = '" + self.growparams + "'\n"
      return self.me

   ## Check if we want to allow growisofs to cross the 4gb boundary
   def is4gbsupported(self):
      processi = popen2.Popen4("uname -s -r")
      status = processi.wait()
      if not os.WIFEXITED(status):
         return 1
      if os.WEXITSTATUS(status) != 0:
         return 1
      strres = processi.fromchild.readline()[0:-1]
      version = re.search(r"Linux (\d+)\.(\d+)\.(\d+)", strres)
      if not version: # Non-Linux: allow
         return 1
      
      if (int(version.group(1)) > 2) or (int(version.group(2)) > 6) or ((int(version.group(1)) == 2) and (int(version.group(2)) == 6) and (int(version.group(3)) >= 8)):
         return 1
      else:
         return 0

   def collect_freespace(self): # Collects current free space
      self.cmd = self.growcmd + " -F " + self.device
      processi = popen2.Popen4(self.cmd)
      status = processi.wait()
      if not os.WIFEXITED(status):
         raise DVDError(0, "growisofs process did not exit correctly.")
      result = processi.fromchild.read()
      if os.WEXITSTATUS(status) != 0:
         if (os.WEXITSTATUS(status) & 0x7F) == errno.ENOSPC:
            # Kludge to force dvd-handler to return a free space of 0
            self.next_session = 1
            self.capacity = 1
            self.freespace_collected = 1
            return
         else:
            raise DVDError(os.WEXITSTATUS(status), "growisofs returned with an error " + result + ". Please check your are using a patched version of dvd+rw-tools.")
      next_sess = re.search(r"\snext_session=(\d+)\s", result, re.MULTILINE)
      capa = re.search(r"\scapacity=(\d+)\s", result, re.MULTILINE)
   
      if next_sess and capa:
         self.next_session = long(next_sess.group(1))
         self.capacity = long(capa.group(1))
         
         # testing cheat (emulate 4GB boundary at 100MB)
         #if self.next_session > 100000000:
         #   self.capacity = self.next_session
      else:
         raise DVDError(0, "Cannot get next_session and capacity from growisofs.\nReturned: " + result)
      
      self.freespace_collected = 1
      return
   
   def collect_mediumtype(self): # Collects current medium type
      self.lasterror = ""
      cmd = self.dvdrwmediainfo + " " + self.device
      processi = popen2.Popen4(cmd)
      status = processi.wait()
      if not os.WIFEXITED(status):
         raise DVDError(0, self.dvdrwmediainfo + " process did not exit correctly.")
      if os.WEXITSTATUS(status) != 0:
         raise DVDError(0, "Cannot get media info from " + self.dvdrwmediainfo)
         return
      result = processi.fromchild.read()
      
      hardware = re.search(r"INQUIRY:\s+(.*)\n", result, re.MULTILINE)
      mediatype = re.search(r"\sMounted Media:\s+([0-9A-F]{2})h, (\S*)\s", result, re.MULTILINE)
      mediamode = re.search(r"\sMounted Media:\s+[0-9A-F]{2}h, \S* (.*)\n", result, re.MULTILINE)
      status = re.search(r"\sDisc status:\s+(.*)\n", result, re.MULTILINE)
      
      if hardware:
         self.hardwaredevice = hardware.group(1)
      
      if mediatype:
         self.disktype = mediatype.group(2)
      else:
         raise DVDError(0, "Media type not found in " + self.dvdrwmediainfo + " output")
      
      if self.disktype == "DVD-RW":
         if mediamode:
            self.diskmode = mediamode.group(1)
         else:
            raise DVDError(0, "Media mode not found for DVD-RW in " + self.dvdrwmediainfo + " output")
      
      if status:
         self.diskstatus = status.group(1)
      else:
         raise DVDError(0, "Disc status not found in " + self.dvdrwmediainfo + " output")

      
      self.mediumtype_collected = 1
      return

   def is_empty(self):
      if not self.freespace_collected:
         self.collect_freespace();
      
      return 0 == self.next_session

   def is_RW(self):
      if not self.mediumtype_collected:
         self.collect_mediumtype();
      return "DVD-RW" == self.disktype or "DVD+RW" == self.disktype or "DVD-RAM" == self.disktype

   def is_plus_RW(self):
      if not self.mediumtype_collected:
         self.collect_mediumtype();
      return "DVD+RW" == self.disktype

   def is_minus_RW(self):
      if not self.mediumtype_collected:
         self.collect_mediumtype();
      return "DVD-RW" == self.disktype
      
   def is_restricted_overwrite(self):
      if not self.mediumtype_collected:
         self.collect_mediumtype();
      return self.diskmode == "Restricted Overwrite"

   def is_blank(self):
      if not self.mediumtype_collected:
         self.collect_mediumtype();
      
      return self.diskstatus == "blank"

   def free(self):
      if not self.freespace_collected:
         self.collect_freespace();
      
      fr = self.capacity-self.next_session-self.margin
      if fr < 0:
         return 0
      else:
         return fr

   def term_handler(self, signum, frame):
      print 'dvd-handler: Signal term_handler called with signal', signum
      if self.pid != 0:
         print "dvd-handler: Sending SIGTERM to pid", self.pid
         os.kill(self.pid, signal.SIGTERM)
         time.sleep(10)
         print "dvd-handler: Sending SIGKILL to pid", self.pid
         os.kill(self.pid, signal.SIGKILL)
         sys.exit(1)

   def write(self, newvol, partfile):
      # Blank DVD+RW when there is no data on it
      if newvol and self.is_plus_RW() and self.is_blank():
         print "DVD+RW looks brand-new, blank it to fix some DVD-writers bugs."
         self.blank()
         print "Done, now writing the part file."
      
      if newvol and self.is_minus_RW() and (not self.is_restricted_overwrite()):
         print "DVD-RW is in " + self.diskmode + " mode, reformating it to Restricted Overwrite"
         self.reformat_minus_RW()
         print "Done, now writing the part file."
      
      cmd = self.growcmd + self.growparams
      if newvol:
         cmd += " -Z "
      else:
         cmd += " -M "
      cmd += self.device + " " + str(partfile)
      print "Running " + cmd
      oldsig = signal.signal(signal.SIGTERM, self.term_handler)
      proc = popen2.Popen4(cmd)
      self.pid = proc.pid
      status = proc.poll()
      while status == -1:
         line = proc.fromchild.readline()
         while len(line) > 0:
            print line,
            line = proc.fromchild.readline()
         time.sleep(1)
         status = proc.poll()
      self.pid = 0
      print
      signal.signal(signal.SIGTERM, oldsig)
      if os.WEXITSTATUS(status) != 0:
         raise DVDError(os.WEXITSTATUS(status), cmd + " exited with status " + str(os.WEXITSTATUS(status)) + ", signal/status " + str(status))

   def prepare(self):
      if not self.is_RW():
         raise DVDError(0, "I won't prepare a non-rewritable medium")
      
      # Blank DVD+RW when there is no data on it
      if self.is_plus_RW() and self.is_blank():
         print "DVD+RW looks brand-new, blank it to fix some DVD-writers bugs."
         self.blank()
         return # It has been completely blanked: Medium is ready to be used by Bacula
      
      if self.is_minus_RW() and (not self.is_restricted_overwrite()):
         print "DVD-RW is in " + self.diskmode + " mode, reformating it to Restricted Overwrite"
         self.reformat_minus_RW()
         return # Reformated: Medium is ready to be used by Bacula
      
      # TODO: Check if /dev/fd/0 and /dev/zero exists, otherwise, run self.blank()
      if not os.path.exists("/dev/fd/0") or not os.path.exists("/dev/zero"):
         print "/dev/fd/0 or /dev/zero doesn't exist, blank the medium completely."
         self.blank()
         return
      
      cmd = self.dd + " if=/dev/zero bs=1024 count=512 | " + self.growcmd + " -Z " + self.device + "=/dev/fd/0"
      print "Running " + cmd
      oldsig = signal.signal(signal.SIGTERM, self.term_handler)
      proc = popen2.Popen4(cmd)
      self.pid = proc.pid
      status = proc.poll() 
      while status == -1:
         line = proc.fromchild.readline()
         while len(line) > 0:
            print line,
            line = proc.fromchild.readline()
         time.sleep(1)
         status = proc.poll()
      self.pid = 0
      print
      signal.signal(signal.SIGTERM, oldsig)
      if os.WEXITSTATUS(status) != 0:
         raise DVDError(os.WEXITSTATUS(status), cmd + " exited with status " + str(os.WEXITSTATUS(status)) + ", signal/status " + str(status))

   def blank(self):
      cmd = self.growcmd + " -Z " + self.device + "=/dev/zero"
      print "Running " + cmd
      oldsig = signal.signal(signal.SIGTERM, self.term_handler)
      proc = popen2.Popen4(cmd)
      self.pid = proc.pid
      status = proc.poll()
      while status == -1:
         line = proc.fromchild.readline()
         while len(line) > 0:
            print line,
            line = proc.fromchild.readline()
         time.sleep(1)
         status = proc.poll()
      self.pid = 0
      print
      signal.signal(signal.SIGTERM, oldsig)
      if os.WEXITSTATUS(status) != 0:
         raise DVDError(os.WEXITSTATUS(status), cmd + " exited with status " + str(os.WEXITSTATUS(status)) + ", signal/status " + str(status))

   def reformat_minus_RW(self):
      cmd = self.dvdrwformat + " -force " + self.device
      print "Running " + cmd
      oldsig = signal.signal(signal.SIGTERM, self.term_handler)
      proc = popen2.Popen4(cmd)
      self.pid = proc.pid
      status = proc.poll()
      while status == -1:
         line = proc.fromchild.readline()
         while len(line) > 0:
            print line,
            line = proc.fromchild.readline()
         time.sleep(1)
         status = proc.poll()
      self.pid = 0
      print
      signal.signal(signal.SIGTERM, oldsig)
      if os.WEXITSTATUS(status) != 0:
         raise DVDError(os.WEXITSTATUS(status), cmd + " exited with status " + str(os.WEXITSTATUS(status)) + ", signal/status " + str(status))

# class disk ends here.

class DVDError(Exception):
   def __init__(self, errno, value):
      self.errno = errno
      self.value = value
      if self.value[-1] == '\n':
         self.value = self.value[0:-1]
   def __str__(self):
      return str(self.value) + " || errno = " + str(self.errno) + " (" + os.strerror(self.errno & 0x7F) + ")"

def usage():
   print "Wrong number of arguments."
   print """
Usage:

dvd-handler DEVICE test
dvd-handler DEVICE free
dvd-handler DEVICE write APPEND FILE
dvd-handler DEVICE blank

where DEVICE is a device name like /dev/sr0 or /dev/dvd.

Operations:
test      Scan the device and report the information found.
           This operation needs no further arguments.
free      Scan the device and report the available space.
write     Write a part file to disk.
           This operation needs two additional arguments.
           The first indicates to append (0) or restart the
           disk (1). The second is the file to write.
prepare   Prepare a DVD+/-RW for being used by Bacula.
           Note: This is only useful if you already have some
           non-Bacula data on a medium, and you want to use
           it with Bacula. Don't run this on blank media, it
           is useless.
"""
   sys.exit(1)

if len(sys.argv) < 3:
   usage()

dvd = disk(sys.argv[1])

if "free" == sys.argv[2]:
   if len(sys.argv) == 3:
      try:
         free = dvd.free()
      except DVDError, e:
         if e.errno != 0:
            print -e.errno
         else:
            print errno.EPIPE
         print str(e)
      else:
         print free
         print "No Error reported."
   else:
      print "Wrong number of arguments for free operation."
      usage()
elif "prepare" == sys.argv[2]:
   if len(sys.argv) == 3:
      try:
         dvd.prepare()
      except DVDError, e:
         print "Error while preparing medium: ", str(e)
         if e.errno != 0:
            sys.exit(e.errno & 0x7F)
         else:
            sys.exit(errno.EPIPE)
      else:
         print "Medium prepared successfully."
   else:
      print "Wrong number of arguments for prepare operation."
      usage()
elif "test" == sys.argv[2]:
   try:
      print str(dvd)
      print "Blank disk: " + str(dvd.is_blank()) + " ReWritable disk: " + str(dvd.is_RW())
      print "Free space: " + str(dvd.free())
   except DVDError, e:
      print "Error while getting informations: ", str(e)
elif "write" == sys.argv[2]:
   if len(sys.argv) == 5:
      try:
         dvd.write(long(sys.argv[3]), sys.argv[4])
      except DVDError, e:
         print "Error while writing part file: ", str(e)
         if e.errno != 0:
            sys.exit(e.errno & 0x7F)
         else:
            sys.exit(errno.EPIPE)
      else:
         print "Part file " + sys.argv[4] + " successfully written to disk."
   else:
      print "Wrong number of arguments for write operation."
      usage()
      sys.exit(1)
else:
   print "No operation - use test, free, prepare or write."
   print "THIS MIGHT BE A CASE OF DEBUGGING BACULA OR AN ERROR!"
sys.exit(0)
