gshred

#!/usr/bin/env python

import re
import os
import sys
import dbus
import shlex
import traceback
import subprocess
import multiprocessing
import Tkinter as tk


def humanize(n):
  # Using 1KB = 1000B, not 1KB = 2**10B.
  units = ['B', 'KB', 'MB', 'GB', 'TB']
  if n < 10**3:
    v = n / float(1.0)
    u = 'B'
  elif n < 10**6:
    v = n / float(10**3)
    u = 'KB'
  elif n < 10**9:
    v = n / float(10**6)
    u = 'MB'
  elif n < 10**12:
    v = n / float(10**9)
    u = 'GB'
  else:
    v = n / float(10**12)
    u = 'TB'
  return '%.1f %s' % (v, u)

class ShredException(Exception):
  pass

class Shred:

  EXE = '/usr/bin/shred'
  TERMINATED = '[TERMINATED]'

  def __init__(self):
    self._exact = True
    self._path = None
    self._zero = False
    self._iterations = 3

  def exact(self, value=True):
    self._exact = value
    return self

  def path(self, value):
    self._path = value
    return self

  def iterations(self, value):
    if value > 0 and value < 10:
      self._iterations = value
    return self

  def __call__(self, output):
    # Make sure that the path exists.
    if not self._path or not os.path.exists(self._path):
      raise ShredException('Could not find path "%s".' % self._path)

    # Build command line.
    command = '%s --force' % Shred.EXE
    command += ' --iterations=%d' % self._iterations
    if self._zero:
      command += ' --zero'
    if self._exact:
      command += ' --exact'
    if output:
      command += ' --verbose'
    command += ' "%s"' % self._path

    # Launch shred.
    proc = subprocess.Popen(shlex.split(command),
                            shell=False,
                            stdin=subprocess.PIPE,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                            bufsize=0)
 
    # Get the verbose logging output (it's coming from stderr).
    while True:
      line = proc.stderr.readline()
      if not line:
        break
      output.put(line)
    
    # Notify termination status.
    proc.wait()
    output.put('%s %d' % (Shred.TERMINATED, proc.returncode)) 

  @staticmethod
  def launch(pool, shred):
    manager = multiprocessing.Manager()
    output = manager.Queue()
    return pool.apply_async(shred, (output,)), manager, output

class App:

  PULSE_PERIOD = 500
  SAFE_PATH = '/lib/live/mount/'

  def __init__(self):
    self._root = tk.Tk()
    self._create_layout()
    self._device = None
    self._shredding = False
    self._output = None
    self._result = None
    self._ticks = 0
    self._reset()
    self._refresh_devices()

  def _create_layout(self):
    # Configure root element.
    self._root.wm_title('gshred')
    self._root.resizable(width=False, height=False)
    self._root.grid_columnconfigure(0, weight=1)
    self._root.grid_rowconfigure(0, weight=1)

    # Create frame for the log.
    self._logframe = tk.Frame(self._root)
    self._logframe.grid(row=1, column=0, sticky='wens')

    # Create menu frame.
    self._menuframe = tk.Frame(self._root)
    self._menuframe.grid(row=0, column=0, sticky='wens')
    self._menuframe.grid_columnconfigure(0, weight=1)

    # Create textbox for the log output. Make it scrollable.
    self._log = tk.Text(self._logframe, height=10, width=120)
    self._scroll = tk.Scrollbar(self._logframe)
    self._scroll.pack(side=tk.RIGHT, fill=tk.Y)
    self._log.pack(side=tk.LEFT, fill=tk.BOTH)
    self._scroll.config(command=self._log.yview)
    self._log.config(yscrollcommand=self._scroll.set)

    # This is a hack to make the 'Text' widget readonly.
    # Source: http://stackoverflow.com/a/34811313
    self._log.bind("<Key>", lambda e: "break")

    # Create dropdown list of all devices (will be populated later).
    self._text = tk.StringVar(self._menuframe, "Select file/device")
    self._devices = tk.OptionMenu(self._menuframe, self._text, tuple())
    self._devices.grid(row=0, column=0, sticky='we')

    # Create shred button.
    self._button = tk.Button(self._menuframe, text="Shred!",
                             command=self._shred)
    self._button.grid(row=0, column=1)

  def _reset(self):
    self._button.config(state='normal')
    self._shredding = False
    self._output = None
    self._result = None
    self._ticks = 0
  
  def _do_pulse(self): 
    if not self._shredding or not self._output or not self._result:
      # Nothing to do.
      return

    if self._ticks * App.PULSE_PERIOD >= 60 * 1000:
      # There has not been any progress (output) for one minute, abort.
      self._reset()
      self._append_log('Operation timed out.')
      return

    # Check if we have something to read out.
    try:
      line = self._output.get(block=False)
    except:
      # Nothing in the queue.
      if self._result.ready():
        # We are done.
        self._append_log('Done!')
        self._reset()
        return
      else:
        # Not done yet, but there is also nothing in the queue.
        # We have to wait a bit.
        self._ticks += 1
        return

    # We have a new message.
    self._ticks = 0
    if line.startswith(Shred.TERMINATED):
      code = int(re.search(r'] (\d+)', line).group(1))
      if code == 0:
        self._append_log('Completed successfully!')
      else:
        self._append_log('Failure.')
      self._reset()
    else:
      self._append_log(line)

  def _pulse(self):
    try:
      self._do_pulse()
    except:
      # Clear log.
      self._log.delete('1.0', tk.END)
      self._append_log(traceback.format_exc())

    # Reschedule pulsing callback.
    self._root.after(App.PULSE_PERIOD, self._pulse)

  def _append_log(self, line):
    self._log.insert(tk.INSERT, line.strip() + '\n')

  def _shred(self):
    try:
      self._do_shred()
    except:
      msg = traceback.format_exc()
      self._append_log(msg)

  def _do_shred(self):
    # Already in progress.
    if self._shredding:
      return

    # Clear log.
    self._log.delete('1.0', tk.END)

    # No device specified.
    device = self._text.get()
    if not device or 'Select file/device' in device:
      self._append_log('No device specified.')
      return

    # Extract the file.
    try:
      device = re.search(r'^(/\S+)', device).group(1)
    except:
      self._append_log('No device specified.')
      return

    # Device does not seem to exist.
    device = os.path.normpath(os.path.abspath(device))
    if not os.path.exists(device):
      self._append_log(('Device or file "%s" does not exist, or you do' +
                       ' not have access.\nTry running as root.') % device)
      return
    
    # Build the shredder.
    shred = Shred().iterations(1).exact(True).path(device)
    pool = multiprocessing.Pool()
    self._append_log('Shredding..')
    self._result, self._manager, self._output = Shred.launch(pool, shred)
    self._button.config(state='disabled')
    self._shredding = True
    
  def _get_devices(self):
    udisk = 'org.freedesktop.UDisks'
    system_bus = dbus.SystemBus()

    udisk_object = system_bus.get_object(udisk, '/org/freedesktop/UDisks')
    udisk_interface = dbus.Interface(udisk_object, udisk)

    # Grab information from all devices.
    devices = list()
    for device in udisk_interface.EnumerateDevices():
      device_object = system_bus.get_object(udisk, device)
      device_interface = dbus.Interface(device_object, dbus.PROPERTIES_IFACE)

      # Grab properties.
      prop = lambda x: device_interface.Get(udisk + '.Device', x)
      mount = prop('DeviceMountPaths')
      devices.append({
        'file': str(prop('DeviceFile')),
        'vendor': str(prop('DriveVendor')),
        'mount': str(mount[0]) if len(mount) > 0 else '',
        'size': humanize(prop('PartitionSize')),
        'model': str(prop('DriveModel')),
      })

    return devices

  def _refresh_devices(self):
    # Clear current list.
    self._text.set('Select file/device')
    self._devices['menu'].delete(0, 'end') 

    # Debugging purposes..
    if len(sys.argv) > 1 and sys.argv[1] == '--debug':
      d = '/root/testfile-small'
      self._devices['menu'].add_command(label=d,
                                        command=tk._setit(self._text, d))
      d = '/root/testfile-big'
      self._devices['menu'].add_command(label=d,
                                        command=tk._setit(self._text, d))

    # Grab devices.
    devices = self._get_devices()
    for device in devices:
      if App.SAFE_PATH in device['mount']:
        # Do not show the live media partitions. :-)
        continue

      if device['size'] == '0.0 B':
        # Do not show special things that do not have a size to them.
        continue

      # Build label.
      label = '%s - %s - %s (%s)' % (device['file'], device['size'],
                                     device['model'], device['vendor'])
      if device['mount']:
        label += ' - %s' % device['mount']

      # Add to dropdown.
      self._devices['menu'].add_command(label=label,
                                        command=tk._setit(self._text, label))
    
  def run(self):
    self._root.after(App.PULSE_PERIOD, self._pulse)
    tk.mainloop()

def main():
  app = App()
  app.run()

if __name__ == '__main__':
  main()