1 files changed, 69 insertions, 0 deletions
diff --git a/get-agent.py b/get-agent.py
new file mode 100755
index 0000000..8f82a23
--- /dev/null
+++ b/get-agent.py
@@ -0,0 +1,69 @@
+#!/usr/bin/env python
+#
+# 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 sys
+import os
+import shutil
+import imp
+import subprocess
+import time
+import re
+import platform
+
+def upgrade(account='Azure', ref='2.0'):
+    #Define variables
+    agentUri = ('https://raw.githubusercontent.com/{0}/'
+                'WALinuxAgent/{1}/waagent').format(account, ref)
+    distro = platform.linux_distribution()
+    cmd = ['service', 'waagent', 'restart']
+    agent_file="/usr/sbin/waagent"
+
+    if "Ubuntu" in distro[0]:
+        cmd[1]='walinuxagent'
+    if "CoreOS" in distro[0]:
+        cmd = ['systemctl', 'restart', 'waagent']
+        agent_file = "/usr/share/oem/bin/waagent"
+
+    if os.path.isfile('waagent'):
+        os.remove('waagent')
+    print "Download WAAgent from: {0}".format(agentUri)
+    try:
+        import urllib2
+        response = urllib2.urlopen(agentUri)
+        html = response.read()
+        with open("waagent", "w+") as F:
+            F.write(html)
+    except:
+        subprocess.call(['wget', agentUri])
+
+    print "Upgrade WAAgent"
+
+    shutil.copyfile("waagent", agent_file)
+    os.chmod(agent_file, 0700)
+
+    job = subprocess.Popen(cmd)
+    job.wait()
+    
+if __name__ == '__main__':
+    if len(sys.argv) == 3:
+        upgrade(sys.argv[1], sys.argv[2])
+    else:
+        upgrade()