Skip to content

Commit c832c49

Browse files
wip
1 parent 2136eb4 commit c832c49

File tree

7 files changed

+210
-87
lines changed

7 files changed

+210
-87
lines changed

DynamoDbEncryption/runtimes/python/src/aws_database_encryption_sdk/encryptor/client.py

Lines changed: 38 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,15 @@
44
from botocore.paginate import Paginator
55
import botocore.client
66
from typing import Callable, Optional
7+
import inspect
78

89
from aws_database_encryption_sdk.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb_transforms.models import (
910
GetItemInputTransformInput,
1011
GetItemOutputTransformInput,
1112
PutItemInputTransformInput,
1213
PutItemOutputTransformInput,
14+
BatchWriteItemInputTransformInput,
15+
BatchWriteItemOutputTransformInput,
1316
)
1417
from aws_database_encryption_sdk.transform import (
1518
dict_to_ddb,
@@ -73,6 +76,13 @@ def _copy_sdk_response_to_dbesdk_response(self, sdk_response, dbesdk_response):
7376
if sdk_response_key not in dbesdk_response:
7477
dbesdk_response[sdk_response_key] = sdk_response_value
7578

79+
def _get_protected_methods(self):
80+
"""Return a list of all protected methods in the given object."""
81+
return [
82+
name for name, member in inspect.getmembers(self, predicate=inspect.ismethod)
83+
if name.startswith('_') and not name.startswith('__')
84+
]
85+
7686
def put_item(self, **kwargs):
7787
# TODO: refactor shared logic (DDB/Python conversions, client/table)
7888
dynamodb_input = self._maybe_transform_request_to_dynamodb_item(item_key = "Item", **kwargs)
@@ -95,8 +105,8 @@ def put_item(self, **kwargs):
95105
)
96106
).transformed_output
97107
self._copy_sdk_response_to_dbesdk_response(sdk_response, dbesdk_response)
98-
self._maybe_transform_response_to_python_dict(dbesdk_response)
99-
return response
108+
self._maybe_transform_response_item_to_python_dict(dbesdk_response)
109+
return dbesdk_response
100110

101111
def get_item(self, **kwargs):
102112
dynamodb_input = self._maybe_transform_request_to_dynamodb_item(item_key = "Key", **kwargs)
@@ -113,11 +123,34 @@ def get_item(self, **kwargs):
113123
)
114124
).transformed_output
115125
self._copy_sdk_response_to_dbesdk_response(sdk_response, dbesdk_response)
116-
self._maybe_transform_response_to_python_dict(dbesdk_response)
117-
return response
126+
self._maybe_transform_response_item_to_python_dict(dbesdk_response)
127+
return dbesdk_response
128+
129+
def batch_write_item(self, **kwargs):
130+
# dynamodb_input = self._maybe_transform_request_to_dynamodb_item(item_key = "Key", **kwargs)
131+
dynamodb_input = kwargs
132+
transformed_request = self._transformer.batch_write_item_input_transform(
133+
BatchWriteItemInputTransformInput(
134+
sdk_input = dynamodb_input
135+
)
136+
).transformed_input
137+
sdk_response = self._client.batch_write_item(**transformed_request)
138+
dbesdk_response = self._transformer.batch_write_item_output_transform(
139+
BatchWriteItemOutputTransformInput(
140+
original_input = dynamodb_input,
141+
sdk_output = sdk_response,
142+
)
143+
).transformed_output
144+
self._copy_sdk_response_to_dbesdk_response(sdk_response, dbesdk_response)
145+
# self._maybe_transform_response_to_python_dict(dbesdk_response)
146+
return dbesdk_response
118147

119148
def __getattr__(self, name):
120149
if hasattr(self._client, name):
121150
print(f'calling underlyign client {name=}')
122151
return getattr(self._client, name)
123-
raise KeyError("idk")
152+
# __getattr__ doesn't find protected methods by default.
153+
elif name in self._get_protected_methods():
154+
return getattr(self, name)
155+
else:
156+
raise KeyError("idk still")

TestVectors/dafny/DDBEncryption/src/DecryptManifest.dfy

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,17 @@ module {:options "-functionSyntax:4"} DecryptManifest {
2020
import opened JSONHelpers
2121
import JsonConfig
2222
import ENC = AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorTypes
23+
import KeyVectors
2324

24-
method OnePositiveTest(name : string, config : JSON, encrypted : JSON, plaintext : JSON) returns (output : Result<bool, string>)
25+
method OnePositiveTest(name : string, config : JSON, encrypted : JSON, plaintext : JSON, keys : KeyVectors.KeyVectorsClient)
26+
returns (output : Result<bool, string>)
27+
requires keys.ValidState()
28+
modifies keys.Modifies
29+
ensures keys.ValidState()
2530
{
2631
var enc :- JsonConfig.GetRecord(encrypted);
2732
var plain :- JsonConfig.GetRecord(plaintext);
28-
var encryptor :- JsonConfig.GetItemEncryptor(name, config);
33+
var encryptor :- JsonConfig.GetItemEncryptor(name, config, keys);
2934
var decrypted :- expect encryptor.DecryptItem(
3035
ENC.DecryptItemInput(
3136
encryptedItem:=enc.item
@@ -36,10 +41,14 @@ module {:options "-functionSyntax:4"} DecryptManifest {
3641
return Success(true);
3742
}
3843

39-
method OneNegativeTest(name : string, config : JSON, encrypted : JSON) returns (output : Result<bool, string>)
44+
method OneNegativeTest(name : string, config : JSON, encrypted : JSON, keys: KeyVectors.KeyVectorsClient)
45+
returns (output : Result<bool, string>)
46+
requires keys.ValidState()
47+
modifies keys.Modifies
48+
ensures keys.ValidState()
4049
{
4150
var enc :- JsonConfig.GetRecord(encrypted);
42-
var encryptor :- JsonConfig.GetItemEncryptor(name, config);
51+
var encryptor :- JsonConfig.GetItemEncryptor(name, config, keys);
4352
var decrypted := encryptor.DecryptItem(
4453
ENC.DecryptItemInput(
4554
encryptedItem:=enc.item
@@ -51,7 +60,11 @@ module {:options "-functionSyntax:4"} DecryptManifest {
5160
return Success(true);
5261
}
5362

54-
method OneTest(name : string, value : JSON) returns (output : Result<bool, string>)
63+
method OneTest(name : string, value : JSON, keys: KeyVectors.KeyVectorsClient)
64+
returns (output : Result<bool, string>)
65+
requires keys.ValidState()
66+
modifies keys.Modifies
67+
ensures keys.ValidState()
5568
{
5669
:- Need(value.Object?, "Test must be an object");
5770

@@ -89,15 +102,19 @@ module {:options "-functionSyntax:4"} DecryptManifest {
89102

90103
if types.value == "positive-decrypt" {
91104
:- Need(plaintext.Some?, "positive-decrypt Test requires a 'plaintext' member.");
92-
output := OnePositiveTest(name, config.value, encrypted.value, plaintext.value);
105+
output := OnePositiveTest(name, config.value, encrypted.value, plaintext.value, keys);
93106
} else if types.value == "negative-decrypt" {
94-
output := OneNegativeTest(name, config.value, encrypted.value);
107+
output := OneNegativeTest(name, config.value, encrypted.value, keys);
95108
} else {
96109
return Failure("Invalid encrypt type : '" + types.value + "'.");
97110
}
98111
}
99112

100-
method Decrypt(inFile : string) returns (output : Result<bool, string>)
113+
method Decrypt(inFile : string, keyVectors: KeyVectors.KeyVectorsClient)
114+
returns (output : Result<bool, string>)
115+
requires keyVectors.ValidState()
116+
modifies keyVectors.Modifies
117+
ensures keyVectors.ValidState()
101118
{
102119
var timeStamp :- expect Time.GetCurrentTimeStamp();
103120
print timeStamp + " Decrypt : ", inFile, "\n";
@@ -154,7 +171,7 @@ module {:options "-functionSyntax:4"} DecryptManifest {
154171
for i := 0 to |tests.value| {
155172
var obj := tests.value[i];
156173
:- Need(obj.1.Object?, "Value of test '" + obj.0 + "' must be an Object.");
157-
var _ :- OneTest(obj.0, obj.1);
174+
var _ :- OneTest(obj.0, obj.1, keyVectors);
158175
}
159176

160177
timeStamp :- expect Time.GetCurrentTimeStamp();

TestVectors/dafny/DDBEncryption/src/EncryptManifest.dfy

Lines changed: 29 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@ module {:options "-functionSyntax:4"} EncryptManifest {
2121
import opened JSONHelpers
2222
import JsonConfig
2323
import ENC = AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorTypes
24+
import KeyVectors
25+
26+
const DEFAULT_KEYS : string := "../../../submodules/MaterialProviders/TestVectorsAwsCryptographicMaterialProviders/dafny/TestVectorsAwsCryptographicMaterialProviders/test/keys.json"
2427

2528
function Manifest() : (string, JSON)
2629
{
@@ -42,10 +45,14 @@ module {:options "-functionSyntax:4"} EncryptManifest {
4245
("client", Object(result))
4346
}
4447

45-
method OnePositiveTest(name : string, theType : string, desc : string, config : JSON, decryptConfig : Option<JSON>, record : JSON) returns (output : Result<(string, JSON), string>)
48+
method OnePositiveTest(name : string, theType : string, desc : string, config : JSON, decryptConfig : Option<JSON>, record : JSON, keys: KeyVectors.KeyVectorsClient)
49+
returns (output : Result<(string, JSON), string>)
50+
requires keys.ValidState()
51+
modifies keys.Modifies
52+
ensures keys.ValidState()
4653
{
4754
var rec :- JsonConfig.GetRecord(record);
48-
var encryptor :- JsonConfig.GetItemEncryptor(name, config);
55+
var encryptor :- JsonConfig.GetItemEncryptor(name, config, keys);
4956
var encrypted :- expect encryptor.EncryptItem(
5057
ENC.EncryptItemInput(
5158
plaintextItem:=rec.item
@@ -64,10 +71,14 @@ module {:options "-functionSyntax:4"} EncryptManifest {
6471
return Success((name, Object(result)));
6572
}
6673

67-
method OneNegativeTest(name : string, config : JSON, record : JSON) returns (output : Result<bool, string>)
74+
method OneNegativeTest(name : string, config : JSON, record : JSON, keys: KeyVectors.KeyVectorsClient)
75+
returns (output : Result<bool, string>)
76+
requires keys.ValidState()
77+
modifies keys.Modifies
78+
ensures keys.ValidState()
6879
{
6980
var rec :- JsonConfig.GetRecord(record);
70-
var encryptor :- JsonConfig.GetItemEncryptor(name, config);
81+
var encryptor :- JsonConfig.GetItemEncryptor(name, config, keys);
7182
var encrypted := encryptor.EncryptItem(
7283
ENC.EncryptItemInput(
7384
plaintextItem:=rec.item
@@ -80,7 +91,11 @@ module {:options "-functionSyntax:4"} EncryptManifest {
8091
}
8192

8293

83-
method OneTest(name : string, value : JSON) returns (output : Result<Option<(string, JSON)>, string>)
94+
method OneTest(name : string, value : JSON, keys: KeyVectors.KeyVectorsClient)
95+
returns (output : Result<Option<(string, JSON)>, string>)
96+
requires keys.ValidState()
97+
modifies keys.Modifies
98+
ensures keys.ValidState()
8499
{
85100
:- Need(value.Object?, "Test must be an object");
86101

@@ -117,20 +132,24 @@ module {:options "-functionSyntax:4"} EncryptManifest {
117132
:- Need(record.Some?, "Test requires a 'record' member.");
118133

119134
if types.value == "positive-encrypt" {
120-
var x :- OnePositiveTest(name, "positive-decrypt", description.value, config.value, decryptConfig, record.value);
135+
var x :- OnePositiveTest(name, "positive-decrypt", description.value, config.value, decryptConfig, record.value, keys);
121136
return Success(Some(x));
122137
} else if types.value == "negative-decrypt" {
123-
var x :- OnePositiveTest(name, "negative-decrypt", description.value, config.value, decryptConfig, record.value);
138+
var x :- OnePositiveTest(name, "negative-decrypt", description.value, config.value, decryptConfig, record.value, keys);
124139
return Success(Some(x));
125140
} else if types.value == "negative-encrypt" {
126-
var _ := OneNegativeTest(name, config.value, record.value);
141+
var _ := OneNegativeTest(name, config.value, record.value, keys);
127142
return Success(None);
128143
} else {
129144
return Failure("Invalid encrypt type : '" + types.value + "'.");
130145
}
131146
}
132147

133-
method Encrypt(inFile : string, outFile : string, lang : string, version : string) returns (output : Result<bool, string>)
148+
method Encrypt(inFile : string, outFile : string, lang : string, version : string, keyVectors: KeyVectors.KeyVectorsClient)
149+
returns (output : Result<bool, string>)
150+
requires keyVectors.ValidState()
151+
modifies keyVectors.Modifies
152+
ensures keyVectors.ValidState()
134153
{
135154
var timeStamp :- expect Time.GetCurrentTimeStamp();
136155
print timeStamp + " Encrypt : ", inFile, "\n";
@@ -187,7 +206,7 @@ module {:options "-functionSyntax:4"} EncryptManifest {
187206
for i := 0 to |tests.value| {
188207
var obj := tests.value[i];
189208
:- Need(obj.1.Object?, "Value of test '" + obj.0 + "' must be an Object.");
190-
var newTest :- OneTest(obj.0, obj.1);
209+
var newTest :- OneTest(obj.0, obj.1, keyVectors);
191210
if newTest.Some? {
192211
test := test + [newTest.value];
193212
}

TestVectors/dafny/DDBEncryption/src/Index.dfy

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,14 @@ module WrappedDDBEncryptionMain {
1414
import FileIO
1515
import JSON.API
1616
import opened JSONHelpers
17+
import KeyVectors
18+
import KeyVectorsTypes = AwsCryptographyMaterialProvidersTestVectorKeysTypes
1719

18-
method AddJson(prev : TestVectorConfig, file : string) returns (output : Result<TestVectorConfig, string>)
20+
method AddJson(prev : TestVectorConfig, file : string, keyVectors: KeyVectors.KeyVectorsClient)
21+
returns (output : Result<TestVectorConfig, string>)
22+
requires keyVectors.ValidState()
23+
modifies keyVectors.Modifies
24+
ensures keyVectors.ValidState()
1925
{
2026
var configBv := FileIO.ReadBytesFromFile(file);
2127
if configBv.Failure? {
@@ -24,20 +30,31 @@ module WrappedDDBEncryptionMain {
2430
}
2531
var configBytes := BvToBytes(configBv.value);
2632
var json :- expect API.Deserialize(configBytes);
27-
output := ParseTestVector(json, prev);
33+
output := ParseTestVector(json, prev, keyVectors);
2834
if output.Failure? {
2935
print output.error, "\n";
3036
}
3137
}
3238

33-
method ASDF() {
39+
method ASDF()
40+
{
41+
// Create a singleton keyVectors client used in every test.
42+
// Right now, all test vectors use the same keys manifest, located at DEFAULT_KEYS.
43+
// Parsing JSON is expensive in some languages.
44+
//
45+
var keyVectors :- expect KeyVectors.KeyVectors(
46+
KeyVectorsTypes.KeyVectorsConfig(
47+
keyManifestPath := DEFAULT_KEYS
48+
)
49+
);
50+
3451
WriteSetPermutations.WriteSetPermutations();
3552
var config := MakeEmptyTestVector();
36-
config :- expect AddJson(config, "records.json");
37-
config :- expect AddJson(config, "configs.json");
38-
config :- expect AddJson(config, "data.json");
39-
config :- expect AddJson(config, "iotest.json");
40-
config :- expect AddJson(config, "PermTest.json");
41-
config.RunAllTests();
53+
config :- expect AddJson(config, "records.json", keyVectors);
54+
config :- expect AddJson(config, "configs.json", keyVectors);
55+
config :- expect AddJson(config, "data.json", keyVectors);
56+
config :- expect AddJson(config, "iotest.json", keyVectors);
57+
config :- expect AddJson(config, "PermTest.json", keyVectors);
58+
config.RunAllTests(keyVectors);
4259
}
4360
}

0 commit comments

Comments
 (0)