Skip to content

Commit 74c89e7

Browse files
Chronicle Teamcopybara-github
Chronicle Team
authored and
copybara-github
committed
Adds sample Python script for UpdateSubject BK API
PiperOrigin-RevId: 403087605
1 parent badf1cb commit 74c89e7

2 files changed

Lines changed: 241 additions & 0 deletions

File tree

‎access_control/update_subject.py‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
#!/usr/bin/env python3
2+
3+
# Copyright 2021 Google LLC
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
#
17+
"""Executable and reusable sample for updating information about a subject.
18+
19+
A subject can be an analyst or an Identity Provider (IdP) group.
20+
"""
21+
22+
import argparse
23+
import json
24+
import sys
25+
from typing import Any, Mapping
26+
from typing import Optional
27+
from typing import Sequence
28+
29+
from google.auth.transport import requests
30+
31+
from common import chronicle_auth
32+
from common import regions
33+
34+
CHRONICLE_API_BASE_URL = "https://backstory.googleapis.com"
35+
36+
37+
def initialize_command_line_args(
38+
args: Optional[Sequence[str]] = None) -> Optional[argparse.Namespace]:
39+
"""Initializes and checks all the command-line arguments."""
40+
parser = argparse.ArgumentParser()
41+
chronicle_auth.add_argument_credentials_file(parser)
42+
regions.add_argument_region(parser)
43+
parser.add_argument(
44+
"-n", "--name", type=str, required=True, help="subject ID")
45+
parser.add_argument(
46+
"-t",
47+
"--type",
48+
type=str,
49+
required=True,
50+
help=("the subject's type (ANALYST or IDP_GROUP)"))
51+
parser.add_argument(
52+
"-rs",
53+
"--roles",
54+
type=str,
55+
required=True,
56+
help=("the role(s) the subject must have after the update"))
57+
58+
# Sanity checks for the command-line arguments.
59+
60+
# No need for a sanity check for the subject name and roles because these
61+
# arguments convert the provided input into strings and accept a wide range
62+
# of values. If the subject name isn't passed in, the error will be thrown
63+
# from the argparse library.
64+
65+
# Check the subject type.
66+
parsed_args = parser.parse_args(args)
67+
if parsed_args.type not in ("ANALYST", "IDP_GROUP"):
68+
print("Error: invalid subject type")
69+
return None
70+
71+
return parser.parse_args(args)
72+
73+
74+
def update_subject(http_session: requests.AuthorizedSession, name: str,
75+
subject_type: str,
76+
roles: Sequence[str]) -> Mapping[str, Sequence[Any]]:
77+
"""Updates information about a subject.
78+
79+
Args:
80+
http_session: Authorized session for HTTP requests.
81+
name: The ID of the subject to retrieve information about.
82+
subject_type: The subject's type (ANALYST or IDP_GROUP).
83+
roles: The role(s) the subject must have after the update.
84+
85+
Returns:
86+
Information about the requested subject in the form:
87+
{
88+
"subject": {
89+
"name": "test@test.com",
90+
"type": "SUBJECT_TYPE_ANALYST",
91+
"roles": [
92+
{
93+
"name": "Test",
94+
"title": "Test role",
95+
"description": "The Test role",
96+
"createTime": "yyyy-mm-ddThh:mm:ss.ssssssZ",
97+
"isDefault": false,
98+
"permissions": [
99+
{
100+
"name": "Test",
101+
"title": "Test permission",
102+
"description": "The Test permission",
103+
"createTime": "yyyy-mm-ddThh:mm:ss.ssssssZ",
104+
},
105+
...
106+
]
107+
},
108+
...
109+
]
110+
}
111+
}
112+
113+
Raises:
114+
requests.exceptions.HTTPError: HTTP request resulted in an error
115+
(response.status_code >= 400).
116+
"""
117+
url = f"{CHRONICLE_API_BASE_URL}/v1/subjects/{name}"
118+
body = {
119+
"name": name,
120+
"type": subject_type,
121+
"roles": roles,
122+
}
123+
update_fields = ["subject.roles"]
124+
params = {"update_mask": ",".join(update_fields)}
125+
response = http_session.request("PATCH", url, params=params, json=body)
126+
127+
if response.status_code >= 400:
128+
print(response.text)
129+
response.raise_for_status()
130+
return response.json()
131+
132+
133+
if __name__ == "__main__":
134+
cli = initialize_command_line_args()
135+
if not cli:
136+
sys.exit(1) # A sanity check failed.
137+
138+
CHRONICLE_API_BASE_URL = regions.url(CHRONICLE_API_BASE_URL, cli.region)
139+
session = chronicle_auth.initialize_http_session(cli.credentials_file)
140+
print(
141+
json.dumps(
142+
update_subject(session, cli.name, cli.type, cli.roles.split(",")),
143+
indent=2))

‎access_control/update_subject_test.py‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
"""Unit tests for the "update_subject" module."""
16+
17+
import unittest
18+
import argparse
19+
20+
from unittest import mock
21+
22+
from google.auth.transport import requests
23+
24+
from . import update_subject
25+
26+
27+
class UpdateSubjectTest(unittest.TestCase):
28+
29+
def test_initialize_command_line_args(self):
30+
actual = update_subject.initialize_command_line_args(
31+
["--name=test@test.com", "--type=ANALYST", "--roles="])
32+
self.assertEqual(
33+
actual,
34+
argparse.Namespace(
35+
credentials_file=None,
36+
name="test@test.com",
37+
type="ANALYST",
38+
roles="",
39+
region="us"))
40+
41+
def test_initialize_command_line_args_invalid_type(self):
42+
actual = update_subject.initialize_command_line_args(
43+
["--name=test@test.com", "--type=INVALID", "--roles="])
44+
self.assertIsNone(actual)
45+
46+
@mock.patch.object(requests, "AuthorizedSession", autospec=True)
47+
@mock.patch.object(requests.requests, "Response", autospec=True)
48+
def test_update_subject_error(self, mock_response, mock_session):
49+
mock_session.request.return_value = mock_response
50+
type(mock_response).status_code = mock.PropertyMock(return_value=400)
51+
mock_response.raise_for_status.side_effect = (
52+
requests.requests.exceptions.HTTPError())
53+
54+
with self.assertRaises(requests.requests.exceptions.HTTPError):
55+
update_subject.update_subject(mock_session, "", "", [])
56+
57+
@mock.patch.object(requests, "AuthorizedSession", autospec=True)
58+
@mock.patch.object(requests.requests, "Response", autospec=True)
59+
def test_update_subject(self, mock_response, mock_session):
60+
mock_session.request.return_value = mock_response
61+
type(mock_response).status_code = mock.PropertyMock(return_value=200)
62+
subject_id = "test@test.com"
63+
subject_type = "ANALYST"
64+
roles = ["Test"]
65+
expected = {
66+
"subject": {
67+
"name":
68+
subject_id,
69+
"type":
70+
"SUBJECT_TYPE_ANALYST",
71+
"roles": [{
72+
"name":
73+
"Test",
74+
"title":
75+
"Test role",
76+
"description":
77+
"The Test role",
78+
"createTime":
79+
"2020-11-05T00:00:00Z",
80+
"isDefault":
81+
False,
82+
"permissions": [{
83+
"name": "Test",
84+
"title": "Test permission",
85+
"description": "The Test permission",
86+
"createTime": "2020-11-05T00:00:00Z",
87+
},]
88+
},]
89+
},
90+
}
91+
mock_response.json.return_value = expected
92+
actual = update_subject.update_subject(mock_session, subject_id,
93+
subject_type, roles)
94+
self.assertEqual(actual, expected)
95+
96+
97+
if __name__ == "__main__":
98+
unittest.main()

0 commit comments

Comments
 (0)