Python script to expand VirtualBox VMDK disk (created to expand vagrant boxes using VB provider)

expand_disk.py

import subprocess,re,os,random,string,sys,fileinput
from sys import platform as _platform
from time import sleep

# VirtualBox disk resize tool
# Author: Geronimo Onativia (geronimox@gmail.com)
# Some ideas taken from
# https://github.com/mitchellh/vagrant/issues/2339#issuecomment-230839878
# http://stackoverflow.com/a/31177761/13116

DEFAULT_DISK = "box-disk1.vmdk"
DEFAULT_VM_NAME = "bravento_api"
DEFAULT_RESIZED_MARK = "_resized"
DEFAULT_DISK_SIZE_MB = 12288
RND = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(5))


def run_command(command):
    p = subprocess.Popen(command,
                         stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE,
                         shell=True)
    # Read stdout from subprocess until the buffer is empty !
    for line in iter(p.stdout.readline, b''):
        if line: # Don't print blank lines
            yield line
    # This ensures the process has completed, AND sets the 'returncode' attr
    while p.poll() is None:                                                                                                                                        
        sleep(.1) #Don't waste CPU-cycles
    # Empty STDERR buffer
    err = p.stderr.read()
    if p.returncode != 0:
       # The run_command() function is responsible for logging STDERR 
       print "Error: " + err  

def touch(fname, times=None):
    with open(fname, 'a'):
        os.utime(fname, times)

def find_in_file(filename, pattern):
    p = re.compile(pattern)
    for i, line in enumerate(open(filename)):
        for match in re.finditer(pattern, line):
            return match.groups()[0]
    return ""

print "Beware that your VM must have a name set via Vagrantfile."
vm_name = raw_input('Enter name given to vm [bravento_api]: ')
if not vm_name:
    vm_name = DEFAULT_VM_NAME
disk_size = raw_input("Enter new disk size in MB ["+str(DEFAULT_DISK_SIZE_MB)+"]: ")
if not disk_size:
    disk_size = DEFAULT_DISK_SIZE_MB

if _platform == "linux" or _platform == "linux2" or _platform == "darwin":
   vb_machine_folder_cmd='VBoxManage list systemproperties | grep "Default machine folder"'
elif _platform == "win32":
   vb_machine_folder_cmd='VBoxManage list systemproperties | findstr "Default\\ machine\\ folder"'
   
output = subprocess.check_output(vb_machine_folder_cmd, shell=True)
vb_machines_folder = re.split(':\s+', output)[1].strip()
vb_machine_resized = os.path.join(vb_machines_folder, vm_name+DEFAULT_RESIZED_MARK)
vb_machine_folder = os.path.join(vb_machines_folder, vm_name)
vb_machine_disk = os.path.join(vb_machine_folder, DEFAULT_DISK)
vb_machine_vbox = os.path.join(vb_machine_folder, vm_name+".vbox")




if os.path.isdir(vb_machine_folder):
    if os.path.exists(vb_machine_disk):
        old_disk_size=os.stat(vb_machine_disk).st_size
        if old_disk_size<(float(disk_size)*1000*1000*0.99) and not os.path.exists(vb_machine_resized):
            print "Resize needed as size is " + str(int(old_disk_size/float(1000*1000)))+"MB < "+str(disk_size)+"MB"
            
            tmp1 = os.path.join(vb_machine_folder, RND+"_tmp1")
            tmp2 = os.path.join(vb_machine_folder, RND+"_tmp2")
            uuid = find_in_file(vb_machine_vbox,"HardDisk uuid=\"{([\w-]+)}\"")
            print "Got current disk UUID "+uuid

            command = "VBoxManage clonemedium disk \"" + vb_machine_disk + "\" \"" + tmp1 + "\" --format vdi"
            print ">"+command
            for line in run_command(command):
                print line
            command = "VBoxManage modifyhd \"" + tmp1 + "\" --resize "+str(disk_size)
            print ">"+command
            for line in run_command(command):
                print line
            command = "VBoxManage clonemedium disk \"" + tmp1 + "\" \"" + tmp2 + "\" --format vmdk"
            print ">"+command
            for line in run_command(command):
                print line

            if os.path.exists(tmp2) and os.path.exists(tmp1):
                os.remove(tmp1)
                os.remove(vb_machine_disk)
                os.rename(tmp2,vb_machine_disk)
                print "Re-setting UUID for the new disk"
                command = "VBoxManage internalcommands sethduuid \"" + vb_machine_disk + "\" "+uuid 
                print ">"+command
                for line in run_command(command):
                    print line
                touch(vb_machine_resized)
                print "VMDK Expanded Successfully!!"
                print "Now run this commands to allow sda1 to take new available space:"
                print "vagrant up --no-provision"
                print "vagrant ssh --command \"sudo resize2fs /dev/sda1\""
                print "vagrant provision"
            else:
                print "Failed to expand to "+str(disk_size)+" the disk "+vb_machine_disk

            if os.path.exists(tmp2):
                os.remove(tmp2)
            if os.path.exists(tmp1):
                os.remove(tmp1)
            
        else:
            print "Resize not needed as size is " + str(int(old_disk_size/float(1000*1000)))+"MB > "+str(disk_size)+"MB or file "+vm_name+DEFAULT_RESIZED_MARK+" already exists at "+vb_machines_folder
    else:
        print "Not a file! > "+vb_machine_disk
else:
    print "Not a dir! > "+vb_machine_folder