Example Scripts

"""
This script is intended to be used as a testing script for private integrations and sends a single 
test user account resource to Vanta as written. You can use this script as a starting point 
to build your own private integration.
"""

import argparse
import datetime
import json
import os
import requests

TOKEN_ENDPOINT = "/oauth/token"


# This function returns a single fake user account resource.
# In a real app this would return a list of all users in your system.
def get_users():
    return [
        {
            "displayName": "Test User",
            "uniqueId": "123",
            "externalUrl": "https://example.com",
            "fullName": "Test User",
            "accountName": "Test User",
            "email": "[email protected]",
            "permissionLevel": "ADMIN",
            "createdTimestamp": datetime.datetime.now().isoformat() + "Z",
            "status": "ACTIVE",
            "mfaEnabled": False,
            "mfaMethods": [],
            "authMethod": "PASSWORD",
        }
    ]


CONFIG = {
    # Retrieve these values from the Developer Console page for your app.
    "client_id": os.environ.get("PRIVATE_INTEGRATION_CLIENT_ID"),
    "client_secret": os.environ.get("PRIVATE_INTEGRATION_CLIENT_ID"),
    "resource_id": os.environ.get(
        "USER_ACCOUNT_RESOURCE_ID",
    ),
    "base_url": "https://api.vanta.com",
}


def get_access_token(config):
    headers = {
        "Content-Type": "application/json",
    }

    data = {
        "client_id": config["client_id"],
        "client_secret": config["client_secret"],
        "scope": "connectors.self:write-resource connectors.self:read-resource",
        "grant_type": "client_credentials",
    }

    r = requests.post(
        f'{config["base_url"]}{TOKEN_ENDPOINT}', headers=headers, data=json.dumps(data)
    )

    if r.status_code != 200:
        raise Exception(
            f"Error getting access token with status code {r.status_code}: {r.text}"
        )

    resp_text = json.loads(r.text)
    return resp_text["access_token"]


# This function syncs a single resource to Vanta.
def send_data(config, access_token):
    headers = {
        "Authorization": f"Bearer {access_token}",
        "Content-Type": "application/json",
    }

    data = {
        "resourceId": config["resource_id"],
        "resources": get_users(), # Replace with your internal system logic for getting users, or any other user.
    }

    # This is the endpoint for syncing the user account resource type. For a different resource type, fetch
    # the appropriate endpoint from https://developer.vanta.com/reference/
    url = f'{config["base_url"]}/v1/resources/user_account/sync_all'
    r = requests.put(url, headers=headers, data=json.dumps(data))

    if r.status_code != 200:
        raise Exception(
            f"Error sending data with status code {r.status_code}: {r.text}"
        )

    else:
        print("Successfully sent data to Vanta.")

# This function lists all resources of a given type sent to Vanta.
def list_data(config, access_token):
    headers = {
        "Authorization": f"Bearer {access_token}",
        "Content-Type": "application/json",
    }

    # This is the endpoint for listing the user account resource type. For a different resource type, fetch
    # the appropriate endpoint from https://developer.vanta.com/reference/
    url = f'{config["base_url"]}/v1/resources/user_account/list_all?resourceId={config["resource_id"]}'
    r = requests.get(url, headers=headers)

    if r.status_code != 200:
        raise Exception(
            f"Error listing data with status code {r.status_code}: {r.text}"
        )

    else:
        print(r.text)


def main(env, action):
    access_token = get_access_token(CONFIG[env])
    if action == "list":
        list_data(CONFIG, access_token)
    else:
        send_data(CONFIG, access_token)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Send test data.")
    parser.add_argument("--action", type=str, required=False, choices=["list", "send"], default="send")
    args = parser.parse_args()
    main(args.env, args.action)