From 2bd830caae6b10dd075f6581dc83874d0a622a52 Mon Sep 17 00:00:00 2001
From: Adam Furbee <adam.furbee@gmail.com>
Date: Wed, 15 Nov 2017 13:11:34 -0600
Subject: [PATCH] added support for extension attributes

---
 extension_attributes/sample_ea.sh            |  2 +
 extension_attributes/templates/sample_ea.xml | 13 +++++
 scripts/templates/sample_script.xml          |  2 +-
 sync.py                                      | 57 +++++++++++++++++++-
 templates/ea.xml                             | 13 +++++
 5 files changed, 84 insertions(+), 3 deletions(-)
 create mode 100755 extension_attributes/sample_ea.sh
 create mode 100644 extension_attributes/templates/sample_ea.xml
 create mode 100644 templates/ea.xml

diff --git a/extension_attributes/sample_ea.sh b/extension_attributes/sample_ea.sh
new file mode 100755
index 0000000..0cc902c
--- /dev/null
+++ b/extension_attributes/sample_ea.sh
@@ -0,0 +1,2 @@
+#!/bin/bash
+echo "<result>hello world</result>"
diff --git a/extension_attributes/templates/sample_ea.xml b/extension_attributes/templates/sample_ea.xml
new file mode 100644
index 0000000..28efef9
--- /dev/null
+++ b/extension_attributes/templates/sample_ea.xml
@@ -0,0 +1,13 @@
+<computer_extension_attribute>
+  <id></id>
+  <name>Sample EA Uploaded by git2jss</name>
+  <description/>
+  <data_type>String</data_type>
+  <input_type>
+    <type>script</type>
+    <platform>Mac</platform>
+    <script></script>
+  </input_type>
+  <inventory_display>Extension Attributes</inventory_display>
+  <recon_display>Extension Attributes</recon_display>
+</computer_extension_attribute>
diff --git a/scripts/templates/sample_script.xml b/scripts/templates/sample_script.xml
index 02baac4..88a31bd 100644
--- a/scripts/templates/sample_script.xml
+++ b/scripts/templates/sample_script.xml
@@ -1,4 +1,4 @@
 <script>
-  <notes>testing testing testing</notes>
+  <name>Sample Script Uploaded by git2jss</name>
   <script_contents> </script_contents>
 </script>
diff --git a/sync.py b/sync.py
index 2ed7c24..513d302 100755
--- a/sync.py
+++ b/sync.py
@@ -21,11 +21,63 @@ async def upload_profiles(session):
 
 
 async def upload_extension_attributes(session, url, user, passwd):
-    pass
+    mypath = dirname(realpath(__file__))
+    ext_attrs = [f.name for f in os.scandir(join(mypath, 'extension_attributes'))
+                 if f.is_file() and f.name.split('.')[-1] in supported_ea_extensions]
+    tasks = []
+    for ea in ext_attrs:
+        task = asyncio.ensure_future(upload_extension_attribute(session, url, user, passwd, ea))
+        tasks.append(task)
+    responses = await asyncio.gather(*tasks)
 
 
 async def upload_extension_attribute(session, url, user, passwd, ext_attr):
-    pass
+    mypath = dirname(realpath(__file__))
+    auth = aiohttp.BasicAuth(user, passwd)
+    headers = {'content-type': 'application/xml'}
+    with open(join(mypath, 'extension_attributes/' + ext_attr), 'r') as f:
+        data=f.read()
+    with async_timeout.timeout(10):
+        template = await get_ea_template(session, url, user, passwd, ext_attr)
+        async with session.get(url + '/JSSResource/computerextensionattributes/name/' + template.find('name').text,
+                auth=auth) as resp:
+            template.find('input_type/script').text = data
+            if args.verbose:
+                print(ET.tostring(template))
+            if resp.status == 200:
+                put_url = url + '/JSSResource/computerextensionattributes/name/' + template.find('name').text
+                resp = await session.put(put_url, auth=auth, data=ET.tostring(template), headers=headers)
+            else:
+                post_url = url + '/JSSResource/computerextensionattributes/id/0'
+                resp = await session.post(post_url, auth=auth, data=ET.tostring(template), headers=headers)
+    if resp.status in (201, 200):
+        print(f'Uploaded Extension Attribute {ext_attr}')
+    return resp.status
+
+
+async def get_ea_template(session, url, user, passwd, ext_attr):
+    auth = aiohttp.BasicAuth(user, passwd)
+    mypath = dirname(realpath(__file__))
+    try:
+        with open(join(mypath, 'extension_attributes/templates/' + ext_attr.split('.')[0] + '.xml'), 'r') as file:
+            template = ET.fromstring(file.read())
+    except FileNotFoundError:
+        with async_timeout.timeout(10):
+            async with session.get(url + '/JSSResource/computerextensionattributes/name/' + ext_attr,
+                    auth=auth) as resp:
+                if resp.status == 200:
+                    async with session.get(url + '/JSSResource/computerextensionattributes/name/' + ext_attr, auth=auth) as response:
+                        template = ET.fromstring(await response.text())
+                else:
+                    template = ET.parse(join(mypath, 'templates/ea.xml')).getroot()
+    # name is mandatory, so we use the filename if nothing is set in a template
+    if args.verbose:
+        print(ET.tostring(template))
+    if template.find('name') is None:
+        ET.SubElement(template, 'name').text = ext_attr
+    elif template.find('name').text is '' or template.find('name').text is None:
+        template.find('name').text = ext_attr
+    return template
 
 
 async def upload_scripts(session, url, user, passwd):
@@ -90,6 +142,7 @@ async def get_script_template(session, url, user, passwd, script):
 async def main(args):
     async with aiohttp.ClientSession() as session:
         await upload_scripts(session, args.url, args.username, args.password)
+        await upload_extension_attributes(session, args.url, args.username, args.password)
 
 
 if __name__ == '__main__':
diff --git a/templates/ea.xml b/templates/ea.xml
new file mode 100644
index 0000000..4eed87b
--- /dev/null
+++ b/templates/ea.xml
@@ -0,0 +1,13 @@
+<computer_extension_attribute>
+  <id></id>
+  <name></name>
+  <description/>
+  <data_type>String</data_type>
+  <input_type>
+    <type>script</type>
+    <platform>Mac</platform>
+    <script></script>
+  </input_type>
+  <inventory_display>Extension Attributes</inventory_display>
+  <recon_display>Extension Attributes</recon_display>
+</computer_extension_attribute>