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.gcp;
18 
19 import com.google.api.gax.paging.AbstractPagedListResponse;
20 import com.google.cloud.kms.v1.CreateCryptoKeyRequest;
21 import com.google.cloud.kms.v1.CreateKeyRingRequest;
22 import com.google.cloud.kms.v1.CryptoKey;
23 import com.google.cloud.kms.v1.CryptoKeyName;
24 import com.google.cloud.kms.v1.CryptoKeyVersion;
25 import com.google.cloud.kms.v1.CryptoKeyVersionName;
26 import com.google.cloud.kms.v1.CryptoKeyVersionTemplate;
27 import com.google.cloud.kms.v1.ImportCryptoKeyVersionRequest;
28 import com.google.cloud.kms.v1.ImportJob;
29 import com.google.cloud.kms.v1.KeyManagementServiceClient;
30 import com.google.cloud.kms.v1.KeyRing;
31 import com.google.cloud.kms.v1.KeyRingName;
32 import com.google.cloud.kms.v1.LocationName;
33 import com.google.cloud.kms.v1.ProtectionLevel;
34 import com.google.protobuf.ByteString;
35 
36 import java.util.Optional;
37 import java.util.stream.Stream;
38 import java.util.stream.StreamSupport;
39 
40 /** GCP client convenience wrapper for interacting with a key ring. */
41 public class KeyRingClient implements AutoCloseable {
42     final KeyRingName mKeyRingName;
43     final KeyManagementServiceClient mClient;
44 
KeyRingClient(KeyRingName keyRingName)45     public KeyRingClient(KeyRingName keyRingName) throws Exception {
46         this.mKeyRingName = keyRingName;
47         this.mClient = KeyManagementServiceClient.create();
48     }
49 
50     /**
51      * Create the key ring corresponding to this client's KeyRingName.
52      *
53      * <p>Should only be run ONCE.
54      */
createKeyRing()55     KeyRing createKeyRing() {
56         return mClient.createKeyRing(
57                 CreateKeyRingRequest.newBuilder()
58                         .setParent(
59                                 LocationName.of(
60                                                 mKeyRingName.getProject(),
61                                                 mKeyRingName.getLocation())
62                                         .toString())
63                         .setKeyRingId(mKeyRingName.getKeyRing())
64                         .build());
65     }
66 
getKeyRing()67     public KeyRing getKeyRing() {
68         return mClient.getKeyRing(mKeyRingName);
69     }
70 
71     /** Create a regular GCP KMS key (non import). */
createCryptoKey(String cryptoKeyId)72     CryptoKey createCryptoKey(String cryptoKeyId) {
73         return mClient.createCryptoKey(mKeyRingName, cryptoKeyId, CryptoKey.getDefaultInstance());
74     }
75 
76     /** Find by name. */
findCryptoKey(String cryptoKeyId)77     Optional<CryptoKey> findCryptoKey(String cryptoKeyId) {
78         String cryptoKeyName =
79                 CryptoKeyName.of(
80                                 mKeyRingName.getProject(),
81                                 mKeyRingName.getLocation(),
82                                 mKeyRingName.getKeyRing(),
83                                 cryptoKeyId)
84                         .toString();
85         return stream(mClient.listCryptoKeys(mKeyRingName))
86                 .filter(cryptoKey -> cryptoKey.getName().equals(cryptoKeyName))
87                 .findFirst();
88     }
89 
90     /** Find the version 1 of a crypto key by name. */
findCryptoKeyVersion(String cryptoKeyId)91     Optional<CryptoKeyVersion> findCryptoKeyVersion(String cryptoKeyId) {
92         return findCryptoKeyVersion(cryptoKeyId, "1");
93     }
94 
95     /** Find a specific crypto key version. */
findCryptoKeyVersion(String cryptoKeyId, String versionId)96     Optional<CryptoKeyVersion> findCryptoKeyVersion(String cryptoKeyId, String versionId) {
97         String cryptoKeyVersionName =
98                 CryptoKeyVersionName.of(
99                                 mKeyRingName.getProject(),
100                                 mKeyRingName.getLocation(),
101                                 mKeyRingName.getKeyRing(),
102                                 cryptoKeyId,
103                                 versionId)
104                         .toString();
105         return findCryptoKey(cryptoKeyId)
106                 .flatMap(
107                         k ->
108                                 stream(mClient.listCryptoKeyVersions(k.getName()))
109                                         .filter(ckv -> ckv.getName().equals(cryptoKeyVersionName))
110                                         .findFirst());
111     }
112 
113     /** Import a local private key to GCP KMS. */
importCryptoKey( CryptoKey cryptoKey, ImportJob importJob, byte[] wrappedKeyMaterial)114     CryptoKeyVersion importCryptoKey(
115             CryptoKey cryptoKey, ImportJob importJob, byte[] wrappedKeyMaterial) throws Exception {
116         return mClient.importCryptoKeyVersion(
117                 ImportCryptoKeyVersionRequest.newBuilder()
118                         .setParent(cryptoKey.getName())
119                         .setImportJob(importJob.getName())
120                         .setAlgorithm(
121                                 CryptoKeyVersion.CryptoKeyVersionAlgorithm
122                                         .RSA_SIGN_PKCS1_2048_SHA256)
123                         .setRsaAesWrappedKey(ByteString.copyFrom(wrappedKeyMaterial))
124                         .build());
125     }
126 
createCryptoKeyForImport( String cryptoKeyId, ProtectionLevel protectionLevel, CryptoKeyVersion.CryptoKeyVersionAlgorithm algorithm)127     CryptoKey createCryptoKeyForImport(
128             String cryptoKeyId,
129             ProtectionLevel protectionLevel,
130             CryptoKeyVersion.CryptoKeyVersionAlgorithm algorithm) {
131         return mClient.createCryptoKey(
132                 CreateCryptoKeyRequest.newBuilder()
133                         .setParent(mKeyRingName.toString())
134                         .setCryptoKeyId(cryptoKeyId)
135                         .setCryptoKey(
136                                 CryptoKey.newBuilder()
137                                         .setPurpose(CryptoKey.CryptoKeyPurpose.ASYMMETRIC_SIGN)
138                                         .setVersionTemplate(
139                                                 CryptoKeyVersionTemplate.newBuilder()
140                                                         .setProtectionLevel(protectionLevel)
141                                                         .setAlgorithm(algorithm)
142                                                         .build())
143                                         .setImportOnly(true))
144                         .setSkipInitialVersionCreation(true)
145                         .build());
146     }
147 
createImportJob(String cryptoKeyId)148     ImportJob createImportJob(String cryptoKeyId) {
149         ImportJob importJob =
150                 mClient.createImportJob(
151                         mKeyRingName,
152                         cryptoKeyId,
153                         ImportJob.newBuilder()
154                                 .setProtectionLevel(ProtectionLevel.SOFTWARE)
155                                 .setImportMethod(ImportJob.ImportMethod.RSA_OAEP_3072_SHA1_AES_256)
156                                 .build());
157 
158         while (mClient.getImportJob(importJob.getName()).getState()
159                 != ImportJob.ImportJobState.ACTIVE) {
160             try {
161                 Thread.sleep(1000);
162             } catch (InterruptedException e) {
163                 throw new RuntimeException(e);
164             }
165         }
166 
167         return importJob;
168     }
169 
170     /** Utility to turn paged responses into streams. */
stream(AbstractPagedListResponse<?, ?, T, ?, ?> response)171     private static <T> Stream<T> stream(AbstractPagedListResponse<?, ?, T, ?, ?> response) {
172         return StreamSupport.stream(response.iterateAll().spliterator(), false);
173     }
174 
175     @Override
close()176     public void close() throws Exception {
177         this.mClient.close();
178     }
179 }
180