Merge pull request #25 from kaincenteno/activity_download

adds a download util to download the resulting file from OrgDeviceActivity object

adds an initial wait time of 10 seconds for the result of assign/unassign devices to prevent rate limitting

updates parameters from device_id to device_ids when refering to a list of devices

updates unassignment/assignment command to include downloading the CSV with the result

updates cli.md and readme.md with changes

Author: kaincenteno

README.md

@@ -79,17 +79,17 @@ device_assigned_server = axm_client.list_devices_in_mdm_server(device_id='SERIAL
 print(device_assigned_server)
 
 assignment_result = axm_client.assign_unassign_device_to_mdm_server(
-    device_id='SERIAL_NUMBER',
+    device_ids=['SERIAL_NUMBER', "ANOTHER_SERIAL_NUMBER"],
     server_id="MDM_SERVER_ID",
     action="ASSIGN_DEVICES"|"UNASSIGN_DEVICES"
 )
+print(assignment_result)
 
 apple_care_coverage = axm_client.get_apple_care_coverage(device_id='SERIAL_NUMBER')
 print(apple_care_coverage)
 ```
 
 ## Issues:
 * need to add tests
-* unassign, assign devices need to be able to pass more than 1 device
 
 This is still a work in progress

docs/cli.md

@@ -130,12 +130,12 @@ Assign one or more devices to an MDM server.
 **Usage**:
 
 ```console
-$ pyaxm-cli assign-device [OPTIONS] DEVICE_ID... SERVER_ID
+$ pyaxm-cli assign-device [OPTIONS] DEVICE_IDS... SERVER_ID
 ```
 
 **Arguments**:
 
-* `DEVICE_ID...`: [required]
+* `DEVICE_IDS...`: [required]
 * `SERVER_ID`: [required]
 
 **Options**:
@@ -149,12 +149,12 @@ Unassign one or more devices from an MDM server.
 **Usage**:
 
 ```console
-$ pyaxm-cli unassign-device [OPTIONS] DEVICE_ID... SERVER_ID
+$ pyaxm-cli unassign-device [OPTIONS] DEVICE_IDS... SERVER_ID
 ```
 
 **Arguments**:
 
-* `DEVICE_ID...`: [required]
+* `DEVICE_IDS...`: [required]
 * `SERVER_ID`: [required]
 
 **Options**:

pyaxm/cli.py

@@ -4,6 +4,7 @@
 from typing_extensions import Annotated
 from typing import List
 from pyaxm.client import Client
+from pyaxm.utils import download_activity_csv
 
 app = typer.Typer()
 
@@ -71,24 +72,32 @@ def mdm_server_assigned(device_id: Annotated[str, typer.Argument()]):
     df.to_csv(sys.stdout, index=False)
 
 @app.command()
-def assign_device(device_id: Annotated[List[str], typer.Argument()], server_id: Annotated[str, typer.Argument()]):
+def assign_device(device_ids: Annotated[List[str], typer.Argument()], server_id: Annotated[str, typer.Argument()]):
     """Assign one or more devices to an MDM server."""
     client = Client()
-    activity = client.assign_unassign_device_to_mdm_server(device_id, server_id, 'ASSIGN_DEVICES')
+    activity = client.assign_unassign_device_to_mdm_server(device_ids, server_id, 'ASSIGN_DEVICES')
     activity_data = {'id': activity.id}
     activity_data.update(activity.attributes.model_dump())
     df = pd.DataFrame([activity_data])
     df.to_csv(sys.stdout, index=False)
 
+    file_path = download_activity_csv(activity)
+    if file_path:
+        typer.echo(f"Report downloaded successfully to: {file_path}")
+
 @app.command()
-def unassign_device(device_id: Annotated[List[str], typer.Argument()], server_id: Annotated[str, typer.Argument()]):
+def unassign_device(device_ids: Annotated[List[str], typer.Argument()], server_id: Annotated[str, typer.Argument()]):
     """Unassign one or more devices from an MDM server."""
     client = Client()
-    activity = client.assign_unassign_device_to_mdm_server(device_id, server_id, 'UNASSIGN_DEVICES')
+    activity = client.assign_unassign_device_to_mdm_server(device_ids, server_id, 'UNASSIGN_DEVICES')
     activity_data = {'id': activity.id}
     activity_data.update(activity.attributes.model_dump())
     df = pd.DataFrame([activity_data])
     df.to_csv(sys.stdout, index=False)
+    
+    file_path = download_activity_csv(activity)
+    if file_path:
+        typer.echo(f"Report downloaded successfully to: {file_path}")
 
 if __name__ == "__main__":
     app()

pyaxm/client.py

@@ -157,10 +157,12 @@ def assign_unassign_device_to_mdm_server(
         )
 
         # use the ID to check the status until it is complete
+        # Waits before the first check to prevent rate limitting
+        time.sleep(10)
         activity_response = self.abm.get_device_activity(unassign_response.data.id, self.access_token.value)
         retry = 0
-        while 'COMPLETED' not in activity_response.data.attributes.status and retry < 5:
-            time.sleep(2 ** retry)
+        while activity_response.data.attributes.status == 'IN_PROGRESS' and retry < 5:
+            time.sleep(10 * retry)
             retry += 1
             activity_response = self.abm.get_device_activity(unassign_response.data.id, self.access_token.value)
 

pyaxm/utils.py

@@ -0,0 +1,39 @@
+import os
+import re
+import requests
+from typing import Optional
+from urllib.parse import urlparse, parse_qs
+from pyaxm.models import OrgDeviceActivity
+
+
+def download_activity_csv(activity: OrgDeviceActivity, save_path: str = None) -> Optional[str]:
+    """
+    Download CSV file from activity if download URL is available.
+    
+    :param activity: OrgDeviceActivity object with attributes.downloadUrl property
+    :param save_path: Optional path to save file (defaults to current directory)
+    :return: Path to saved file or None if no download URL
+    """
+    if not save_path:
+        save_path = os.getcwd()
+
+    if not (activity.attributes and activity.attributes.downloadUrl):
+        return None
+
+    url = activity.attributes.downloadUrl
+    parsed_url = urlparse(url)
+    query_params = parse_qs(parsed_url.query)
+    
+    disposition = query_params['response-content-disposition'][0]
+    match = re.search(r'filename="([^"]+)"', disposition)
+    filename = match.group(1)
+    
+    file_path = os.path.join(save_path, filename)
+    
+    response = requests.get(url)
+    response.raise_for_status()
+    
+    with open(file_path, 'wb') as f:
+        f.write(response.content)
+    
+    return file_path