summaryrefslogtreecommitdiff
path: root/tests/azure_test.py
blob: fd152b7df188ec38eaabeb9db872582adcbb8043 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
#!/usr/bin/env python

import os
import sys
import imp
import json
import time
import pwd

# waagent has no '.py' therefore create waagent module import manually.
waagent=imp.load_source('waagent','waagent')

from waagent import RunGetOutput,  Run, LoggerInit

"""
Test waagent in azure using azure-cli
Usage:

./azure_test.py --stable_vm_image b4590d9e3ed742e4a1d46e5424aa335e__openSUSE-12.3-v120 --source_disk "http://mystorage.blob.core.windows.net/vhds/my-suse2.vhd" --acct "<storage acct key>" --testname my-osuse --mount_point /mnt/disk --agent_path ../waagent --stable_vm_acct_name MyUserName --stable_vm_acct_pass 'myp455wd' --test_vm_acct_name MyUserName --test_vm_acct_pass 'myp455wd' --azure_location "East US" --part_num 1 --retries 20 --fstype scsi --test_vm_acct_cert /root/.ssh/myCert.pem --stable_vm_acct_cert /root/.ssh/myCert.pem --keep_test_vm_vhd no --teardown_test_vm always --prompt no

azure_test --vm <stable vm name> --testname <testname> [--acct <storage account>]
[--disk <vhd url to use as initial disk image>]
If --disk is specified, use this vhd as starting point,
otherwise use a copy of the stable vm as the starting vhd.

Starting VHD is attached to stable vm and the sources are copied to it.
Spin up a new VM using the VHD.
Loop waiting for provisioned.
If not provisioned:
    Destroy vm and attach the disk to the stable vm.
    Copy the logs to the localhost.
    Destroy all created objects except the starting vhd.
    Exit(1)
If Provosioned:
    Copy the logs to the local host.
    Exit(0)
    
EXAMPLE:

sudo ./azure_test.py --vm my-stable-vm-name --disk "http://mystorageaccount.blob.core.windows.net/myvhds/my-new-os.vhd" --acct mylong-unquoted-starage-account-id --testname my-vm-test --mount_point /my-stablevm-mountpoint --agent_path ../waagent --vm_acct_name root --testvm_acct_name azureuser --testvm_acct_pass 'azureuserpassword' --location "East US" --part_num 2 --retry 20 --fstype bsd --testvm_acct_cert /home/azureuser/.ssh/myCert.pem --keep_vhd "once" --teardown "always" --prompt "no"   
"""

def makeDiskImage(di_name,vhd_url,location,copy=False):
    """
    Create new data disk image of the VHD.  If 'copy'
    is set to True, then create a new VHD in the form of myvhd-di.vhd
    based on the VHD source name.  If 'copy is set to False, re-use the
    vhd.  Returns return code and diskimageVHD path.  Code is 0 on
    success or the azure-cli error code upon error.
    """
    if copy :
        target = os.path.dirname(vhd_url)
        target = target+ '/' + di_name +  '.vhd'
    else :
        target = vhd_url
    cmd='azure vm disk create --json ' + di_name + ' --blob-url ' + target +  ' ' + vhd_url
    print cmd
    waagent.Log( cmd)
    code,output=RunGetOutput(cmd,False)
    print output,code
    waagent.Log(output)
    waagent.Log(str(code))
    return target,code

def makeVMImage(vmi_name,vhd_url,copy=False):
    """
    Create new VM Image based on Disk Image.
    Returns 0 on success or error code upon error.
    """
    if copy :
        target = os.path.dirname(vhd_url)
        target = target+ '/' + vmi_name +  '.vhd'
    else :
        target = vhd_url
    cmd='azure vm image create --json ' + vmi_name + ' --base-vhd '  + target + ' --os Linux --blob-url ' + vhd_url
    print cmd
    waagent.Log( cmd)
    code,output=RunGetOutput(cmd,False)
    print output,code
    waagent.Log(str(code))
    waagent.Log(output)
    return code

def makeVM(vm_name,vmi_name,vhd_url,test_name,location,vmuser_name,vmuser_pass,vmuser_cert,copy=False):
    """
    Create new VM from the VM Image.
    Returns 0 on success or error code upon error.
    """
    target=os.path.dirname(vhd_url)
    target = target + '/' + test_name +  '.vhd'
    cmd='azure vm create --json '
    if copy :
        cmd += ' --blob-url "' + target + '"'
    else :
        target=vhd_url
    cmd += ' --location "' + location + '"'
    if os.path.exists(vmuser_cert):
        cmd += ' -t "' + vmuser_cert + '"'
    cmd += ' -e 22 ' + vm_name + ' ' + vmi_name + ' ' + vmuser_name + ' \'' +vmuser_pass + '\''
    print cmd
    waagent.Log( cmd)
    code,output=RunGetOutput(cmd,False)
    print output,code
    waagent.Log(str(code))
    waagent.Log(output)
    retry=3
    while code !=0 and retry > 0 :
        time.sleep(5)
        code,output=RunGetOutput(cmd,False)
        retry -=1
    return target

def flushDiskImage(di_name,dele=True):
    """
    Delete the VM Image.
    On error we asume the VM disk image is deleted
    """
    cmd='azure vm disk delete --json '
    if dele :
        cmd += '--blob-delete '
    cmd+= di_name
    print cmd
    waagent.Log( cmd)
    code,output=RunGetOutput(cmd,False)
    print output,code
    waagent.Log( str(code))
    waagent.Log(output)
    return output,code

def flushVMImage(vmi_name):
    """
    Delete the VM Image.
    Always delete the underlying blob.
    On error we asume the VM image is deleted.
    """
    cmd='azure vm image delete --blob-delete --json ' + vmi_name
    print cmd
    waagent.Log( cmd)
    code,output=RunGetOutput(cmd,False)
    print output,code
    waagent.Log( str(code))
    waagent.Log(output)
    return output,code

def flushVM(vm_name,dele=False):
    """
    Delete the VM.
    On error we asume the VM is deleted
    """
    cmd='azure vm delete --json '
    if dele :
        cmd += ' --blob-delete '
    cmd += vm_name
    print cmd
    waagent.Log( cmd)
    code,output=RunGetOutput(cmd,False)
    print output,code
    waagent.Log( str(code)) 
    waagent.Log(output)
    return output,code


def createStableVMFromVMImage(vm_image,vhd_url):
    """
    Create a new stable vm, provisioned with acct and certificate from
    the VMImage, using the basepath of vhd_url for the new VM's vhd.
    """
    stableVM=testname+'-stable-vm'
    return makeVM(stableVM,vm_image,vhd_url,testname+'-stable',location,stableVMaccount,stableVMpass,stableVMCert,copy=True)
    
def createStableVMFromVHD(vmi_name,vhd_url):
    """
    Create a new stable vm, provisioned with acct and certificate from
    the VHD, using the basepath of vhd_url for the new VM's vhd.
    """
    makeVMImage(vmi_name,vhd_url,False)
    return createStableVMFromVMImage(vmi_name,vhd_url)

def createDiskImageFromStableVMDisk(vm_name):
    """
    Determine the media link for the os disk of the stable vm.
    Create vhd disk image copy.  <vm_name>-<testname>-di
    Return new disk_image_media_path or None on error.
    """
    cmd='azure vm disk list --json'
    print cmd
    waagent.Log( cmd)
    code,output=RunGetOutput(cmd,False)
    print output,code
    waagent.Log( str(code))
    waagent.Log(output)
    if code:
        print 'Error is ' + str(code)
        waagent.Log( 'Error is ' + str(code))
        return None
    j=json.loads(output)
    source_media_link=None
    for i in j :
        if i.has_key('AttachedTo'):
            if i['AttachedTo']['RoleName'] == vm_name:
                source_media_link=i['MediaLink']
                break
    if not source_media_link:
        print 'Unable to locate OS disk for ' + vm_name 
        waagent.Log( 'Unable to locate OS disk for ' + vm_name )
        return None
    target_name= testname + '-di'
    makeDiskImage(target_name,source_media_link,location,copy=True)
    target_media_link=os.path.dirname(source_media_link) + '/' + target_name + '.vhd'
    return target_media_link

def addDiskImageToVM(vm_name,di_name):
    """
    Attach the disk image to the 'stableVM'.
    Returns the LUN if successful otherwise returns None
    NOTE: azure vm show may return json matching the disk image
    name yet missing the LUN.  When this occurs, the LUN is '0'. 
    """
    cmd='azure vm disk attach --json ' + vm_name + ' ' + di_name
    print cmd
    waagent.Log( cmd)
    code,output=RunGetOutput(cmd,False)
    print output,code
    waagent.Log(str(code))
    waagent.Log(output)
    cmd='azure vm show --json ' + vm_name
    print cmd
    waagent.Log( cmd)
    code,output=RunGetOutput(cmd,False)
    print output,code
    waagent.Log( str(code))
    waagent.Log(output)
    retries=3
    while code != 0 and retries: 
        retries-=1
        print cmd
        waagent.Log( cmd)
        code,output=RunGetOutput(cmd,False)
        
    if code == 0:
        jsn=json.loads(output)
        for i in jsn['DataDisks']:
            if i['DiskName'] == di_name:
                if 'Lun' in i :
                    return i['Lun']
                else :
                    return u'0'
    return None

    
def dropDiskImageFromVM(vm_name,lun):
    """
    Detach the disk image from the 'stableVM'.
    On Error we assume the disk is no longer attached.
    """
    cmd='azure vm disk detach --json ' + vm_name + ' ' + lun
    print cmd
    waagent.Log( cmd)
    code,output=RunGetOutput(cmd,False)
    print output,code
    waagent.Log( str(code))
    waagent.Log(output)
    return output,code

def checkVMProvisioned(vm_name):
    cmd='azure vm show --json ' + vm_name
    print cmd
    waagent.Log( cmd)
    code,output=RunGetOutput(cmd,False)
    print output,code
    waagent.Log( str(code))
    waagent.Log(output)
    if code ==0 :
        j=json.loads(output)
        print vm_name+' instance status: ', j['InstanceStatus']
        waagent.Log( vm_name+' instance status: ' + j['InstanceStatus'])
        if j['InstanceStatus'] == 'ReadyRole':
            return True, j['InstanceStatus']
    else :
        print 'Error: ' + output , code
        waagent.Log( 'Error: ' + output +  str(code))
    return False, j['InstanceStatus']
        
def updateAgent(agent_path,vm_name,account,cert,disk_mountpoint,mnt_opts,lun,partnum,provisioned_account):
    """
    Copy the agent specified in 'agent' to the Disk
    using the 'stableVM'.
    """
    retries=30
    retry=0
    cmd='uptime'
    while  ssh_command(vm_name,account,cmd)[1] != 0 and  retry < retries :
        time.sleep(10)
        retry+=1
    #setup sudo NOPASSWD
    pss=stableVMpass.replace('$','\$')
    cmd='echo \'' + pss  + '\' > /home/' + account + '/pswd '
    ssh_command(vm_name,account,cmd)
    cmd='echo \'#!/bin/bash\ncat /home/' + account + '/pswd\n\' > /home/' + account + '/pw.sh'
    ssh_command(vm_name,account,cmd)
    cmd='chmod +x /home/' + account + '/pw.sh'
    ssh_command(vm_name,account,cmd)
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A mkdir -p ' + disk_mountpoint
    ssh_command(vm_name,account,cmd)
    retries=3
    # TODO retires here for the mount
    #mount
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A mount ' +mnt_opts + ' ' + lunToDiskName(lun,partnum) + ' ' +disk_mountpoint
    waagent.Log( cmd)
    retry=0
    while  ssh_command(vm_name,account,cmd)[1] not in (0,32) and  retry < retries :
        if retry == 0:
            if 'bsd' in fstype:
                fcmd = "export SUDO_ASKPASS=./pw.sh && sudo -A fsck_ffs -y "
            else :
                fcmd = "export SUDO_ASKPASS=./pw.sh && sudo -A fsck -y "
            fcmd += lunToDiskName(lun,partnum)
            ssh_command(vm_name,account,fcmd)
        time.sleep(2)
        retry+=1

    # remove packaged agent service if present.
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A chroot '+ disk_mountpoint+' dpkg -r walinuxagent'
    ssh_command(vm_name,account,cmd)    # remove Ubuntu walinuxagent agent service if present.
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A rm '+ disk_mountpoint+'/etc/default/walinuxagent'
    ssh_command(vm_name,account,cmd)    # remove Ubuntu walinuxagent agent service if present.
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A chroot '+ disk_mountpoint+' rpm -e WALinuxAgent'
    ssh_command(vm_name,account,cmd)
    #copy agent
    remote_path='/tmp'
    print 'scp ' + agent_path + ' to ' + vm_name + ' ' + account + ':' + remote_path
    waagent.Log( 'scp ' + agent_path + ' to ' + vm_name + ' ' + account + ':' + remote_path)
    retry=0
    while scp_to_host_command(account,vm_name,remote_path,agent_path)[1] != 0 and retry < retries :
        time.sleep(2)
        retry+=1
    # move agent to /usr/sbin
    cmd= 'export SUDO_ASKPASS=./pw.sh && sudo -A cp ' + remote_path +'/waagent '+ disk_mountpoint+'/usr/sbin/waagent'
    ssh_command(vm_name,account,cmd)
    cmd= 'export SUDO_ASKPASS=./pw.sh && sudo -A chmod 755 '+ disk_mountpoint+'/usr/sbin/waagent'
    ssh_command(vm_name,account,cmd)
    # Fix the password file
    if 'bsd' in fstype:
        cmd='export SUDO_ASKPASS=./pw.sh && sudo -A cp /etc/master.passwd ' + disk_mountpoint + '/etc/master.passwd'
        ssh_command(vm_name,account,cmd)
    else :
        cmd='export SUDO_ASKPASS=./pw.sh && sudo -A cp /etc/passwd ' + disk_mountpoint + '/etc/passwd'
        ssh_command(vm_name,account,cmd)
        cmd='export SUDO_ASKPASS=./pw.sh && sudo -A cp /etc/shadow ' + disk_mountpoint + '/etc/shadow'
        ssh_command(vm_name,account,cmd)
    #remove /var/lib/waagent
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A rm -rf ' + disk_mountpoint + '/var/lib/waagent'
    ssh_command(vm_name,account,cmd)
    #remove /var/log/waagent*
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A rm -rf ' + disk_mountpoint + '/var/log/waagent*'
    ssh_command(vm_name,account,cmd)
    #delete the provisioning user
    if 'bsd' in fstype:
        cmd='export SUDO_ASKPASS=./pw.sh && sudo -A chroot '+ disk_mountpoint+' rmuser -y ' + provisioned_account
        ssh_command(vm_name,account,cmd)
    else :
        cmd='export SUDO_ASKPASS=./pw.sh && sudo -A chroot '+ disk_mountpoint+' userdel -f ' + provisioned_account
        ssh_command(vm_name,account,cmd)
        cmd='export SUDO_ASKPASS=./pw.sh && sudo -A chroot '+ disk_mountpoint+' groupdel ' + provisioned_account
        ssh_command(vm_name,account,cmd)
        cmd='export SUDO_ASKPASS=./pw.sh && sudo -A rm -rf ' + disk_mountpoint + '/home/' + provisioned_account 
        ssh_command(vm_name,account,cmd)
    # install agent
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A chroot  '+ disk_mountpoint+'  /usr/sbin/waagent verbose install '
    ssh_command(vm_name,account,cmd)
    cmd="export SUDO_ASKPASS=./pw.sh && sudo -A  sed -i 's/Verbose=n/Verbose=y/' " + disk_mountpoint+"/etc/waagent.conf"
    ssh_command(vm_name,account,cmd)
    #umount
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A umount ' + lunToDiskName(lun,partnum)
    ssh_command(vm_name,account,cmd)
    
def gatherAgentInfo(localpath,vm_name,account,cert,disk_mountpoint,mnt_opts,lun,partnum):
    """
    Copy the /var/lib/waagent, and /var/log directories to
    localhost:localpath.
    """
    retries=30
    retry=0
    cmd='uptime'
    while  ssh_command(vm_name,account,cmd)[1] != 0 and  retry < retries :
        time.sleep(10)
        retry+=1
    #mount
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A mount ' +mnt_opts + ' '  + lunToDiskName(lun,partnum) + ' ' +disk_mountpoint
    print cmd
    waagent.Log( cmd)
    ssh_command(vm_name,account,cmd)
    #copy info
    Run("mkdir -p "+ localpath)
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A mkdir -p /tmp/results'
    ssh_command(vm_name,account,cmd)
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A cp -r  ' + disk_mountpoint + '/var/log /tmp/results/'
    ssh_command(vm_name,account,cmd)
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A cp -r  ' + disk_mountpoint + '/var/lib/waagent /tmp/results/'
    ssh_command(vm_name,account,cmd)
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A chmod -R 777 /tmp/results'
    ssh_command(vm_name,account,cmd)
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A chown -R ' + account + '  /tmp/results'
    ssh_command(vm_name,account,cmd)
    scp_from_host_command(account,vm_name,'/tmp/results/*',localpath)
    #umount
    cmd='export SUDO_ASKPASS=./pw.sh && sudo -A umount ' + lunToDiskName(lun,partnum)
    print cmd
    waagent.Log( cmd)
    ssh_command(vm_name,account,cmd)

def lunToDiskName(lun,partnum):
    if 'bsd' in fstype :
        return lunToFreeBSDDiskName(lun,partnum)
    else :
        return lunToScsiDiskName(lun,partnum)

def lunToScsiDiskName(lun,partnum):
    """
    Convert lun to '/dev/sd[chr(ord('c')+lun)]partnum'
    """
    return str('/dev/sd'+chr( (ord('c')+int(lun)) ) +str(partnum))

def lunToFreeBSDDiskName(lun,partnum):
    """
    Convert lun to '/dev/da' + str(lun) + 'p' + partnum
    """
    return '/dev/da'+ str(int(lun))  + 'p' + str(partnum)

def ssh_command(host,account,cmd):
    """
    Wrapper for an ssh operation.
    """
    if stableVMCert == None:
        if not os.path.exists('./pw.sh'):
            with open('./pw.sh','w') as F:
                F.write('#!/bin/bash\ncat ./pswd\n')
            os.system('chmod +x ./pw.sh')
            with open('./pswd','w') as F:
                F.write(stableVMpass)
        req = "export SSH_ASKPASS=./pw.sh && setsid ssh -T -o StrictHostKeyChecking='no' " + account + "@" + host.lower() + ".cloudapp.net \"" + cmd + "\""
    else :
        req = "ssh -t -o StrictHostKeyChecking='no' " + account + "@" + host.lower() + ".cloudapp.net \"" + cmd + "\""
    print req
    waagent.Log(req)
    code,output=RunGetOutput(req,False)
    print output,code
    waagent.Log(str(code))
    waagent.Log(output.encode('ascii','ignore'))
    return output,code

def scp_to_host_command(account,host,remote_path,local_path):
    """
    Wrapper for an scp operation.  Always uses -r.
    Requires key authentication configured.
    """
    req="scp -o StrictHostKeyChecking='no' -r " + local_path + " " + account + "@" + host.lower() + ".cloudapp.net:" + remote_path
    print req
    waagent.Log( req)
    code,output=RunGetOutput(req,False)
    print output,code
    waagent.Log( str(code))
    waagent.Log(output)
    return output,code

def scp_from_host_command(account,host,remote_path,local_path):
    """
    Wrapper for an scp operation.  Always uses -r.
    Requires key authentication configured.
    """
    req="scp -r " + account + "@" + host.lower() + ".cloudapp.net:" + remote_path + " " + local_path
    print req
    waagent.Log( req)
    code,output=RunGetOutput(req,False)
    print output,code
    waagent.Log( str(code))
    waagent.Log(output)
    return output,code

def teardown(name):
    diskImageName=os.path.splitext(os.path.basename(sourceVHD))[0]+'-di'
    while makeDiskImage(diskImageName,sourceVHD,location,True)[1] !=0 :
        time.sleep(20)
    lun=addDiskImageToVM(stableVM,diskImageName)
    while lun == None :
        time.sleep(2)
        lun=addDiskImageToVM(stableVM,diskImageName)
    out,code=flushVM(vmName,True)
    if code != 0 :
        vmDisk=out[out.find('disk with name ')+len('disk with name '):out.find(' is currently in use')]
        while flushDiskImage(vmDisk,True)[1] != 0 :
            time.sleep(5)
    out,code=flushVMImage(vmImageName)
    if code != 0 :
        vmDisk=out[out.find('disk with name ')+len('disk with name '):out.find(' is currently in use')]
        while flushDiskImage(vmDisk,True)[1] != 0 :
            time.sleep(5)
    gatherAgentInfo(localInfo+'/'+name,stableVM,stableVMaccount,stableVMCert,stableVMMountpoint,mountOptions,lun,partNum)
    print 'Logs for ' + vmName + ' copied to ' + localInfo + '/' + name
    waagent.Log( 'Logs for ' + vmName + ' copied to ' + localInfo + '/' + name)
    # detach and delete the disk image
    while dropDiskImageFromVM(stableVM,lun)[1] != 0 :
        time.sleep(2)
    while flushDiskImage(diskImageName,('no' in keep_vhd))[1] != 0 :
        time.sleep(2)
    out,code=flushVM(stableVM,True)
    if code != 0 :
        stableVMDisk=out[out.find('disk with name ')+len('disk with name '):out.find(' is currently in use')]
        while flushDiskImage(stableVMDisk,True)[1] != 0 :
            time.sleep(5)
    if stableVMImageName:
        while flushVMImage(stableVMImageName,True)[1] != 0 :
            time.sleep(2)
    
def doPrompt() :
    if prompt in ('yes','on'):
        raw_input('Press enter to continue:')

if __name__ == '__main__' :
    """
    Create a disk image and attach it to StableVM.
    Copy the current sources to it.
    Detach and delete the disk image container
    Delete the vm image container and the VM.
    Create a vm image container, and the VM.
    Check the VM for provision succedded:
        if so, exit
        if provisioning failed:
            delete the vm, delete the vm image, create a disk image
            from the VM's vhd, attach the disk to the stable VM, copy
            /var/log and /var/lib/waagent to the localhost.
    """
    stableVM=''
    mountOptions=''
    stableVMMountpoint='/mnt/disk' # need to ensure this is created if not existing.
    sourceVHD=''
    location=""
    localAgent='/home/ericg/Public/git_repos/private/WALinuxAgent-Private/waagent_freebsd'
    localInfo='./logs'
    stableVMaccount='' # Need to ensure root permissions for this to work
    stableVMpass=''
    stableVMCert=''
    provisionedVMaccount=''
    provisionedVMpass=''
    provisionedVMCert=''
    account=''
    testname='azuretest'
    partNum=1
    provision_retries=1
    fstype='scsi'
    keep_vhd='no'
    teardown_test_vm='fail'
    teardown_stable_vm='fail'
    prompt = 'yes'
    stable_vm_vhd=None
    stableVMImageName=None
    stable_vm_image=None
    logfile='azure_test.log'
    """
    We need to create a disk image container and attach it to a stable
    vm in order to copy the current sources to it.  Then we detach it,
    delete the disk image container, create a vm image container, and
    the VM. 
    """

    for i in range(len(sys.argv)) :
        if '--stable_vm' == sys.argv[i] : stableVM=sys.argv[i+1]
        elif '--source_disk' == sys.argv[i]: sourceVHD=sys.argv[i+1]
        elif '--storage_acct' == sys.argv[i]: account=sys.argv[i+1]
        elif '--testname' == sys.argv[i] : testname=sys.argv[i+1]
        elif '--stable_vm_mount_point' == sys.argv[i] : stableVMMountpoint=sys.argv[i+1]
        elif '--agent_path' == sys.argv[i] : localAgent=sys.argv[i+1]
        elif '--stable_vm_acct_name' == sys.argv[i] : stableVMaccount=sys.argv[i+1]
        elif '--stable_vm_acct_pass' == sys.argv[i] : stableVMpass=sys.argv[i+1]
        elif '--stable_vm_acct_cert' == sys.argv[i] : stableVMCert=sys.argv[i+1]
        elif '--test_vm_acct_name' == sys.argv[i] : provisionedVMaccount=sys.argv[i+1]
        elif '--test_vm_acct_pass' == sys.argv[i] : provisionedVMpass=sys.argv[i+1]
        elif '--test_vm_acct_cert' == sys.argv[i] : provisionedVMCert=sys.argv[i+1]
        elif '--azure_location' == sys.argv[i] : location=sys.argv[i+1]
        elif '--mount_opts' == sys.argv[i] : mountOptions=sys.argv[i+1]
        elif '--part_num' == sys.argv[i] : partNum=sys.argv[i+1]
        elif '--retries' == sys.argv[i] : provision_retries=int(sys.argv[i+1])
        elif '--fs_type' == sys.argv[i] : fstype=sys.argv[i+1]
        elif '--keep_test_vm_vhd' == sys.argv[i] : keep_vhd=sys.argv[i+1]
        elif '--teardown_test_vm' == sys.argv[i] : teardown_test_vm=sys.argv[i+1]
        elif '--teardown_stable_vm' == sys.argv[i] : teardown_stable_vm=sys.argv[i+1]
        elif '--prompt' == sys.argv[i] : prompt=sys.argv[i+1]
        elif '--stable_vm_image' == sys.argv[i] : stable_vm_image=sys.argv[i+1]
        elif '--stable_vm_vhd' == sys.argv[i] : stable_vm_vhd=sys.argv[i+1]
        elif '--logfile' == sys.argv[i] : logfile=sys.argv[i+1]
        
    LoggerInit(logfile,'')
    waagent.Log("User: "+ pwd.getpwuid(os.geteuid()).pw_name +" Running Command :\n" + reduce(lambda x, y: x+' '+y,sys.argv))

    if len(stableVM) == 0 and not ( stable_vm_image or stable_vm_vhd ):
        print '--vm <stable vm> must be provided unless --stable_vm_image or --stable_vm_vhd'
        waagent.Log( '--vm <stable vm> must be provided!')
        sys.exit(1)
    else:
        if stable_vm_image:
            sourceVHD=createStableVMFromVMImage(stable_vm_image,sourceVHD)
            stableVM=testname+'-stable-vm'
        elif stable_vm_vhd:
            stableVMImageName=testname+'-stable-vi'
            sourceVHD=createStableVMFromVHD(stableVMImageName,stable_vm_vhd)
            stableVM=testname+'-stable-vm'
        p = False
        retries = provision_retries
        while not p and retries > 0:
            p,out = checkVMProvisioned(stableVM)
            if not p:
                if 'Failed' in out or 'Timeout' in out :
                    break
                print  stableVM + ' Not Provisioned - sleeping on retry:' + str( provision_retries - retries ) 
                waagent.Log(  stableVM + ' Not Provisioned - sleeping on retry:' + str( provision_retries - retries ) )
                time.sleep(30)
                retries -= 1
            else :
                print stableVM + ' Provision SUCCEEDED.'
                waagent.Log( stableVM + ' Provision SUCCEEDED.')
    # done creating the stable vm        
    vmImageName=os.path.splitext(os.path.basename(sourceVHD))[0]+'-vi'
    #flushVMImage(vmImageName)

    # if no disk image name is provided we want to clone the stable vm disk.
    if not sourceVHD:
        sourceVHD=createDiskImageFromStableVMDisk(stableVM)
        if not sourceVHD:
            print 'Errors - unable to create disk image - assuming created'
            waagent.Log( 'Errors - unable to create disk image - assuming created')
        diskImageName=os.path.splitext(os.path.basename(sourceVHD))[0]
    else :
        diskImageName=os.path.splitext(os.path.basename(sourceVHD))[0]+'-di'
        diskImageVHD,code=makeDiskImage(diskImageName,sourceVHD,location,True)
        if code:
            print 'Error - unable to make ' + diskImageName
            waagent.Log( 'Error - unable to make ' + diskImageName)

    lun=addDiskImageToVM(stableVM,diskImageName)
    while lun == None :
        time.sleep(2)
        lun=addDiskImageToVM(stableVM,diskImageName)

    doPrompt()
    updateAgent(localAgent,stableVM,stableVMaccount,stableVMCert,stableVMMountpoint,mountOptions,lun,partNum,provisionedVMaccount)
    doPrompt()
    #reboot to prevent stale mount bugs
    cmd= 'export SUDO_ASKPASS=./pw.sh && sudo -A reboot'
    ssh_command(stableVM,stableVMaccount,cmd)

    while dropDiskImageFromVM(stableVM,lun)[1] != 0 :
        time.sleep(2)
    while flushDiskImage(diskImageName,False)[1] != 0 :
        time.sleep(2)
    vmImageName=os.path.splitext(os.path.basename(sourceVHD))[0]+'-vi'
    flushVMImage(vmImageName)
    vmName=testname+'-vm'
    flushVM(vmName)
    makeVMImage(vmImageName,diskImageVHD,True)
    sourceVHD=makeVM(vmName,vmImageName,sourceVHD,testname,location,provisionedVMaccount,provisionedVMpass,provisionedVMCert,True)
    print 'The new source vhd is ' + sourceVHD
    waagent.Log( 'The new source vhd is ' + sourceVHD)
    p = False
    retries = provision_retries
    while not p and retries > 0:
        p,out = checkVMProvisioned(vmName)
        if not p:
            if 'Failed' in out or 'Timeout' in out :
                break
            print  vmName + ' Not Provisioned - sleeping on retry:' + str( provision_retries - retries ) 
            waagent.Log(  vmName + ' Not Provisioned - sleeping on retry:' + str( provision_retries - retries ) )
            time.sleep(30)
        else :
            print vmName + ' Provision SUCCEEDED.'
            waagent.Log( vmName + ' Provision SUCCEEDED.')
            doPrompt()
            if teardown_test_vm in ('success','always'):
                teardown(testname+'_pass')
            sys.exit(0)
        retries -= 1
    
    print vmName + ' Provision FAILED.'
    waagent.Log( vmName + ' Provision FAILED.')
    doPrompt()
    if teardown_test_vm in ('fail','always'):
                teardown(testname+'_fail')
    sys.exit(1)