From bcbbf38d9f46ff9e09705cd09c491af72f293381 Mon Sep 17 00:00:00 2001
From: Yue Zhang <yuezha@microsoft.com>
Date: Wed, 10 Dec 2014 00:08:44 +0800
Subject: [PATCH] Add repartition logic

---
 get-agent.py | 41 +++++++++++++++++++++++++++++++++++++++++
 waagent      | 13 +++++++------
 2 files changed, 48 insertions(+), 6 deletions(-)
 create mode 100755 get-agent.py

diff --git a/get-agent.py b/get-agent.py
new file mode 100755
index 0000000..44243cc
--- /dev/null
+++ b/get-agent.py
@@ -0,0 +1,41 @@
+#!/usr/bin/env python
+#
+# Windows Azure Linux Agent
+#
+# Copyright 2014 Microsoft Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Implements parts of RFC 2131, 1541, 1497 and
+# http://msdn.microsoft.com/en-us/library/cc227282%28PROT.10%29.aspx
+# http://msdn.microsoft.com/en-us/library/cc227259%28PROT.13%29.aspx
+#
+
+import urllib2
+import shutil
+import imp
+
+if __name__ == '__main__':
+    account = 'yuezh'
+    agentUri = ('https://raw.githubusercontent.com/{0}/'
+                'WALinuxAgent/2.0/waagent').format(account)
+    response = urllib2.urlopen(agentUri)
+    html = response.read()
+    with open("waagent") as F:
+        F.write(html)
+    os.chmod("waagent", 544)
+    shutil.copyfile("waagent", "/usr/sbin/waagent")
+    waagent=imp.load_source('waagent','/usr/sbin/waagent')
+    waagent.MyDistro=waagent.GetMyDistro()
+    waagent.MyDistro.stopAgentService()
+    waagent.MyDistro.startAgentService()
diff --git a/waagent b/waagent
index 58a9201..8d896e6 100644
--- a/waagent
+++ b/waagent
@@ -478,14 +478,15 @@ class AbstractDistro(object):
 
             #Check partition type
             ret = RunGetOutput("sfdisk -q -c " + device + " 1", chk_err=False)[1].rstrip()
-            
             partition = device + "1"
-            if "GPT" in ret and fs != "ntfs":
-                #GPT(Guid Partition Table) is used. 
+            if ("GPT" in ret or "ee" == ret) and fs != "ntfs":
+                #In Ubuntu the previous command returns 'ee'
+                #GPT(Guid Partition Table) is used.
                 #The 1st partition is Windows reserved partition.
-                #Use the 2nd partition as resource disk
-                partition = device + "2"
-                Run("mkfs." + fs + " " + partition)
+                Run("parted {0} rm 1".format(device))
+                Run("parted {0} rm 2".format(device))
+                Run("parted {0} mkpart primary 0% 100%".format(device))
+                Run("mkfs." + fs + " " + partition + " -F")
             elif ret == "7" and fs != "ntfs":
                 Run("sfdisk -c " + device + " 1 83")
                 Run("mkfs." + fs + " " + partition)