1 /*
2  * Copyright (C) 2024 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.apksig.kms.aws;
18 
19 import software.amazon.awssdk.core.SdkBytes;
20 import software.amazon.awssdk.http.urlconnection.UrlConnectionHttpClient;
21 import software.amazon.awssdk.services.kms.KmsClient;
22 import software.amazon.awssdk.services.kms.model.AlgorithmSpec;
23 import software.amazon.awssdk.services.kms.model.AliasListEntry;
24 import software.amazon.awssdk.services.kms.model.CreateAliasRequest;
25 import software.amazon.awssdk.services.kms.model.CreateKeyRequest;
26 import software.amazon.awssdk.services.kms.model.DescribeKeyRequest;
27 import software.amazon.awssdk.services.kms.model.ExpirationModelType;
28 import software.amazon.awssdk.services.kms.model.GetParametersForImportRequest;
29 import software.amazon.awssdk.services.kms.model.GetParametersForImportResponse;
30 import software.amazon.awssdk.services.kms.model.ImportKeyMaterialRequest;
31 import software.amazon.awssdk.services.kms.model.ImportKeyMaterialResponse;
32 import software.amazon.awssdk.services.kms.model.KeyMetadata;
33 import software.amazon.awssdk.services.kms.model.KeySpec;
34 import software.amazon.awssdk.services.kms.model.KeyUsageType;
35 import software.amazon.awssdk.services.kms.model.ListAliasesRequest;
36 import software.amazon.awssdk.services.kms.model.ListAliasesResponse;
37 import software.amazon.awssdk.services.kms.model.NotFoundException;
38 import software.amazon.awssdk.services.kms.model.OriginType;
39 import software.amazon.awssdk.services.kms.model.WrappingKeySpec;
40 
41 import java.util.ArrayList;
42 import java.util.List;
43 import java.util.Optional;
44 
45 public class KeyAliasClient implements AutoCloseable {
46     private static final String ALIAS_PREFIX = "alias/";
47     private final KmsClient mClient;
48 
KeyAliasClient()49     public KeyAliasClient() {
50         mClient = KmsClient.builder().httpClientBuilder(UrlConnectionHttpClient.builder()).build();
51     }
52 
getKeyForAlias(String keyAlias)53     Optional<KeyMetadata> getKeyForAlias(String keyAlias) {
54         try {
55             return Optional.of(
56                     mClient.describeKey(
57                                     DescribeKeyRequest.builder()
58                                             .keyId(ALIAS_PREFIX + keyAlias)
59                                             .build())
60                             .keyMetadata());
61         } catch (NotFoundException _unused) {
62             System.out.println("Requested key alias " + keyAlias + "was not found!");
63             return Optional.empty();
64         }
65     }
66 
67     /** List all key aliases (for test accounts only - pages through all aliases). */
listKeyAliases()68     public List<AliasListEntry> listKeyAliases() {
69         ListAliasesResponse response = mClient.listAliases();
70         List<AliasListEntry> aliases = new ArrayList<>(response.aliases());
71         while (response.truncated()) {
72             response =
73                     mClient.listAliases(
74                             ListAliasesRequest.builder().marker(response.nextMarker()).build());
75             aliases.addAll(response.aliases());
76         }
77         return aliases;
78     }
79 
findKeyAlias(String keyAlias)80     Optional<AliasListEntry> findKeyAlias(String keyAlias) {
81         return listKeyAliases().stream()
82                 .filter(as -> as.aliasName().equals(ALIAS_PREFIX + keyAlias))
83                 .findFirst();
84     }
85 
createKey(String keyAlias, KeySpec keySpec)86     KeyMetadata createKey(String keyAlias, KeySpec keySpec) {
87         KeyMetadata createdKey =
88                 mClient.createKey(
89                                 CreateKeyRequest.builder()
90                                         .keyUsage(KeyUsageType.SIGN_VERIFY)
91                                         .keySpec(keySpec)
92                                         .build())
93                         .keyMetadata();
94 
95         mClient.createAlias(
96                 CreateAliasRequest.builder()
97                         .aliasName(keyAlias)
98                         .targetKeyId(createdKey.keyId())
99                         .build());
100 
101         return createdKey;
102     }
103 
createKeyForImport(String keyAlias, KeySpec keySpec)104     KeyMetadata createKeyForImport(String keyAlias, KeySpec keySpec) {
105         KeyMetadata createdKey =
106                 mClient.createKey(
107                                 CreateKeyRequest.builder()
108                                         .keyUsage(KeyUsageType.SIGN_VERIFY)
109                                         .keySpec(keySpec)
110                                         .origin(OriginType.EXTERNAL)
111                                         .build())
112                         .keyMetadata();
113 
114         mClient.createAlias(
115                 CreateAliasRequest.builder()
116                         .aliasName(ALIAS_PREFIX + keyAlias)
117                         .targetKeyId(createdKey.keyId())
118                         .build());
119 
120         return createdKey;
121     }
122 
getParametersForImport( WrappingKeySpec wrappingKeySpec, AlgorithmSpec wrappingAlgorithm, String keyId)123     GetParametersForImportResponse getParametersForImport(
124             WrappingKeySpec wrappingKeySpec, AlgorithmSpec wrappingAlgorithm, String keyId) {
125         return mClient.getParametersForImport(
126                 GetParametersForImportRequest.builder()
127                         .wrappingKeySpec(wrappingKeySpec)
128                         .wrappingAlgorithm(wrappingAlgorithm)
129                         .keyId(keyId)
130                         .build());
131     }
132 
importKeyMaterial( String keyId, SdkBytes importToken, byte[] wrappedKey)133     ImportKeyMaterialResponse importKeyMaterial(
134             String keyId, SdkBytes importToken, byte[] wrappedKey) {
135         return mClient.importKeyMaterial(
136                 ImportKeyMaterialRequest.builder()
137                         .keyId(keyId)
138                         .expirationModel(ExpirationModelType.KEY_MATERIAL_DOES_NOT_EXPIRE)
139                         .importToken(importToken)
140                         .encryptedKeyMaterial(SdkBytes.fromByteArray(wrappedKey))
141                         .build());
142     }
143 
144     @Override
close()145     public void close() throws Exception {
146         mClient.close();
147     }
148 }
149