#! /usr/bin/env python

import sys
import subprocess



# 1. fdisk -lu , and use output to identify all problematic drives...
p = subprocess.Popen(['fdisk', '-lu'], stdout=subprocess.PIPE)
fdisk = p.communicate()[0].split('\n')

problems = []

letter = 'wibble'
for line in fdisk:
  if line.startswith('Disk /dev/sd'):
    letter = line[12]
  
  ct = 'cylinders, total'
  if ('heads' in line) and ('sectors/track,' in line) and (ct in line) and line.endswith('sectors'):
    start = line.find(ct) + len(ct)
    sectors = int(line[start:-8])
  
  if line.startswith('/dev/sd%s' % letter):
    parts = line.split()
    if parts[1]=='*':
      parts = parts[:1] + parts[2:]

    drive = parts[0]
    start = int(parts[1])
    end = int(parts[2])
    
    if end >= sectors:
      print 'Bad drive detected: %s' % drive
      problems.append({'drive' : drive, 'start' : start, 'end' : end, 'sectors' : sectors})

if len(problems)==0:
  print 'No problems found - exiting'
  sys.exit(0)



# 2. df, to verify that none of the bad drives are mounted...
p = subprocess.Popen(['df'], stdout=subprocess.PIPE)
df = p.communicate()[0].split('\n')

for line in df:
  for problem in problems:
    if line.startswith(problem['drive']):
      print '%s is mounted - emergency stop.' % problem['drive']
      sys.exit(1)



# Loop and fix each problem in turn...
for problem in problems:
  # Get partition table of disk...
  p = subprocess.Popen(['sfdisk', '-d', problem['drive'][:-1]], stdout=subprocess.PIPE)
  current = p.communicate()[0].split('\n')
  
  new_size = problem['sectors'] - problem['start']
  print 'Correcting %s to have size %i' % (problem['drive'], new_size)
  
  replacement = []
  for line in current:
    if line.startswith(problem['drive']):
      print '::Before:', line
      start_str = ', size='
      start = line.find(start_str) + len(start_str)
      end = line.find(', Id=')
      if start==-1 or end==-1:
        raise RuntimeError('Cant parse partition spec line')
      replacement.append('%s%i%s' % (line[:start], new_size, line[end:]))

    else:
      replacement.append(line)

    print replacement[-1]
  
  # Replacement partition table (dangerous bit!)...
  p = subprocess.Popen(['sfdisk', '--force', problem['drive'][:-1]], stdin=subprocess.PIPE)
  p.communicate('\n'.join(replacement))
  
  # Finally, correct the partitions super block...
  subprocess.call(['e2fsck', '-f', problem['drive']])
  subprocess.call(['resize2fs', problem['drive']])