MAAP DAAC acces using Assume Role in Python

AssumeRole.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1bcb95eb-6e37-4030-a7af-0415846ea953",
   "metadata": {},
   "source": [
    "# Direct Access to DAAC Buckets\n",
    "\n",
    "This notebook demonstrates using assume-role to switch to a role that has permission to read directly from several DAACs. This removes the need to the s3credentials endpoint to get 1 hour temporary credential. \n",
    "\n",
    "> Note: the assumerole itself has an expiration, but it's much longer than 1 hour. This is still better though since assume_role does not generate requests to the DAAC account, so it's less likely to cause overload."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9ea520bd-899b-45b2-86e6-95cb6088a3cc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from maap.maap import MAAP\n",
    "maap = MAAP(maap_host='api.maap-project.org')\n",
    "\n",
    "import boto3\n",
    "import rasterio as rio\n",
    "from rasterio.session import AWSSession"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "fa741b0b-dd79-4a19-a36c-dbc20ed65232",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "arn = \"arn:aws:iam::884094767067:role/maap-data-reader\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "4c408a0d-c6ff-4c81-9d78-c538614d0e69",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def switch_role(arn):\n",
    "    print(boto3.client('sts').get_caller_identity().get('Arn'))\n",
    "    response = boto3.client('sts').assume_role(\n",
    "        RoleArn = arn,\n",
    "        RoleSessionName = 'DAAC_direct'\n",
    "    )\n",
    "    creds = response.get('Credentials')\n",
    "    print(creds)\n",
    "    boto3_session = boto3.Session(\n",
    "            aws_access_key_id=creds['AccessKeyId'],\n",
    "            aws_secret_access_key=creds['SecretAccessKey'],\n",
    "            aws_session_token=creds['SessionToken']\n",
    "    )\n",
    "    print(boto3_session.client('sts').get_caller_identity().get('Arn'))\n",
    "    # return an s3 session with AssumeRole\n",
    "    \n",
    "    return boto3_session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7634f2b7-ad17-4c71-8284-db0d4336500f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "arn:aws:sts::884094767067:assumed-role/MAAP-ADE-K8S/i-0bd53611d0f5e90ae\n",
      "{'AccessKeyId': 'ASIA43WBHR7NYERJEFRZ', 'SecretAccessKey': 'K2tdMw1Scw29rcluG9vEZOBR7AuUq/B5j2NdG2vC', 'SessionToken': 'FwoGZXIvYXdzEOn//////////wEaDL1Jjw/h6akaqJvogSKvASg5OJ1MqFfvTOs/5mORkaWpPay/UygExYzEIzJsCGwS07dHVoPRoHZtokNLuxTkzYxKJFhpNGizzYZnIVzQjbueYYGaxgGKVnViSg1Q1WNqfiH/zxeSBgF6EkHFlHF55fe0R0B0EiHWeDeZEKS/WPQLa+dEwq9vCyuz+XvqRFz0QqiAGBHcWk/4F3yW8/sjktjJRUMBVwsVX/oYsqmJKCFI4aG0hw6tME6f0Xldy5IouNvupwYyLdIXUKZD8dZGMh3O/mOdpWEG7c8l3UcxrwPBzu0LDEXN85L4SaRTqSjqZGlpwQ==', 'Expiration': datetime.datetime(2023, 9, 9, 0, 26, 48, tzinfo=tzutc())}\n",
      "arn:aws:sts::884094767067:assumed-role/maap-data-reader/DAAC_direct\n"
     ]
    }
   ],
   "source": [
    "s3=switch_role(arn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7c7567e5-7530-4f9c-aa1a-70b16f93db91",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# ORNL was a bad test as assume_role wasn't required\n",
    "granule_path = \"s3://ornl-cumulus-prod-protected/gedi/GEDI_L4B_Gridded_Biomass/data/GEDI04_B_MW019MW138_02_002_05_R01000M_PS.tif\"\n",
    "granule_path = \"s3://gesdisc-cumulus-prod-protected/Landslide/Global_Landslide_Exposure_Maps.1.0/Pop_exp_stdev.tif\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "8848a39d-b69d-48ae-b5b8-ecca73254f46",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'driver': 'GTiff', 'dtype': 'float32', 'nodata': 0.0, 'width': 43201, 'height': 15841, 'count': 1, 'crs': CRS.from_epsg(4326), 'transform': Affine(0.00833333333333333, 0.0, -180.0,\n",
      "       0.0, -0.00833333333333333, 72.00006000333326), 'blockysize': 1, 'tiled': False, 'interleave': 'band'}\n"
     ]
    }
   ],
   "source": [
    "# Create a Rasterio session from the S3 session\n",
    "aws_session = AWSSession(s3) #I don't know which session was used?\n",
    "\n",
    "with rio.Env(aws_session):\n",
    "    with rio.open(granule_path) as src:\n",
    "        print(src.profile)\n",
    "        test = src.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "dbde2e16-920f-4034-a52c-c5c67a4ced2f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "arn:aws:sts::884094767067:assumed-role/MAAP-ADE-K8S/i-0bd53611d0f5e90ae\n"
     ]
    }
   ],
   "source": [
    "print(boto3.client('sts').get_caller_identity().get('Arn'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "73818dc7-fb57-4ea0-9a90-b2e4865463df",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 15841, 43201)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test.shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}