44import com .nimbusds .jose .JOSEException ;
55import io .opentdf .platform .sdk .*;
66import io .opentdf .platform .sdk .TDF ;
7+ import io .opentdf .platform .sdk .Config .AssertionVerificationKeys ;
78
89import com .google .gson .Gson ;
910import org .apache .commons .codec .DecoderException ;
11+ import org .bouncycastle .crypto .RuntimeCryptoException ;
12+
1013import picocli .CommandLine ;
1114import picocli .CommandLine .HelpCommand ;
1215import picocli .CommandLine .Option ;
2225import java .io .PrintWriter ;
2326import java .nio .ByteBuffer ;
2427import java .nio .channels .FileChannel ;
28+ import java .nio .charset .StandardCharsets ;
29+ import java .nio .file .Files ;
2530import java .nio .file .Path ;
31+ import java .nio .file .Paths ;
2632import java .nio .file .StandardOpenOption ;
2733import java .security .InvalidAlgorithmParameterException ;
2834import java .security .InvalidKeyException ;
2935import java .security .NoSuchAlgorithmException ;
36+ import java .security .spec .InvalidKeySpecException ;
37+ import java .security .spec .PKCS8EncodedKeySpec ;
38+ import java .security .spec .X509EncodedKeySpec ;
39+ import java .security .KeyFactory ;
40+ import java .security .PrivateKey ;
3041import java .text .ParseException ;
3142import java .util .ArrayList ;
43+ import java .util .Base64 ;
3244import java .util .List ;
45+ import java .util .Map ;
3346import java .util .Optional ;
3447import java .util .concurrent .ExecutionException ;
3548import java .util .function .Consumer ;
3952
4053import javax .net .ssl .TrustManager ;
4154
55+
4256@ CommandLine .Command (name = "tdf" , subcommands = {HelpCommand .class })
4357class Command {
4458
59+ private static final String PRIVATE_KEY_HEADER = "-----BEGIN PRIVATE KEY-----" ;
60+ private static final String PRIVATE_KEY_FOOTER = "-----END PRIVATE KEY-----" ;
61+ private static final String PEM_HEADER = "-----BEGIN (.*)-----" ;
62+ private static final String PEM_FOOTER = "-----END (.*)-----" ;
63+
4564 @ Option (names = { "--client-secret" }, required = true )
4665 private String clientSecret ;
4766
@@ -57,6 +76,68 @@ class Command {
5776 @ Option (names = { "-p" , "--platform-endpoint" }, required = true )
5877 private String platformEndpoint ;
5978
79+ private Object correctKeyType (AssertionConfig .AssertionKeyAlg alg , Object key , boolean publicKey ) throws RuntimeException {
80+ if (alg == AssertionConfig .AssertionKeyAlg .HS256 ) {
81+ if (key instanceof String ) {
82+ key = ((String ) key ).getBytes (StandardCharsets .UTF_8 );
83+ return key ;
84+ } else if (key instanceof byte []) {
85+ return key ;
86+ } else {
87+ throw new RuntimeException ("Unexpected type for assertion key" );
88+ }
89+ } else if (alg == AssertionConfig .AssertionKeyAlg .RS256 ) {
90+ if (!(key instanceof String )) {
91+ throw new RuntimeException ("Unexpected type for assertion key" );
92+ }
93+ String pem = (String ) key ;
94+ String pemWithNewlines = pem .replace ("\\ n" , "\n " );
95+ if (publicKey ){
96+ String base64EncodedPem = pemWithNewlines
97+ .replaceAll (PEM_HEADER , "" )
98+ .replaceAll (PEM_FOOTER , "" )
99+ .replaceAll ("\\ s" , "" )
100+ .replaceAll ("\r \n " , "" )
101+ .replaceAll ("\n " , "" )
102+ .trim ();
103+ byte [] decoded = Base64 .getDecoder ().decode (base64EncodedPem );
104+ X509EncodedKeySpec spec = new X509EncodedKeySpec (decoded );
105+ KeyFactory kf = null ;
106+ try {
107+ kf = KeyFactory .getInstance ("RSA" );
108+ } catch (NoSuchAlgorithmException e ) {
109+ throw new RuntimeException (e );
110+ }
111+ try {
112+ return kf .generatePublic (spec );
113+ } catch (InvalidKeySpecException e ) {
114+ throw new RuntimeException (e );
115+ }
116+ }else {
117+ String privateKeyPEM = pemWithNewlines
118+ .replace (PRIVATE_KEY_HEADER , "" )
119+ .replace (PRIVATE_KEY_FOOTER , "" )
120+ .replaceAll ("\\ s" , "" ); // remove whitespaces
121+
122+ byte [] decoded = Base64 .getDecoder ().decode (privateKeyPEM );
123+
124+ PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec (decoded );
125+ KeyFactory kf = null ;
126+ try {
127+ kf = KeyFactory .getInstance ("RSA" );
128+ } catch (NoSuchAlgorithmException e ) {
129+ throw new RuntimeException (e );
130+ }
131+ try {
132+ return kf .generatePrivate (spec );
133+ } catch (InvalidKeySpecException e ) {
134+ throw new RuntimeException (e );
135+ }
136+ }
137+ }
138+ return null ;
139+ }
140+
60141 @ CommandLine .Command (name = "encrypt" )
61142 void encrypt (
62143 @ Option (names = { "-f" , "--file" }, defaultValue = Option .NULL_VALUE ) Optional <File > file ,
@@ -92,9 +173,29 @@ void encrypt(
92173 try {
93174 assertionConfigs = gson .fromJson (assertionConfig , AssertionConfig [].class );
94175 } catch (JsonSyntaxException e ) {
95- throw new RuntimeException ("Failed to parse assertion, expects an list of assertions" , e );
176+ // try it as a file path
177+ try {
178+ String fielJson = new String (Files .readAllBytes (Paths .get (assertionConfig )));
179+ assertionConfigs = gson .fromJson (fielJson , AssertionConfig [].class );
180+ } catch (JsonSyntaxException e2 ) {
181+ throw new RuntimeException ("Failed to parse assertion from file, expects an list of assertions" , e2 );
182+ } catch (Exception e3 ) {
183+ throw new RuntimeException ("Could not parse assertion as json string or path to file" , e3 );
184+ }
185+ }
186+ // iterate through the assertions and correct the key types
187+ for (int i = 0 ; i < assertionConfigs .length ; i ++) {
188+ AssertionConfig config = assertionConfigs [i ];
189+ if (config .signingKey != null && config .signingKey .isDefined ()) {
190+ try {
191+ Object correctedKey = correctKeyType (config .signingKey .alg , config .signingKey .key , false );
192+ config .signingKey .key = correctedKey ;
193+ } catch (Exception e ) {
194+ throw new RuntimeException ("Error with assertion signing key: " + e .getMessage (), e );
195+ }
196+ }
197+ assertionConfigs [i ] = config ;
96198 }
97-
98199 configs .add (Config .withAssertionConfig (assertionConfigs ));
99200 }
100201
@@ -126,15 +227,50 @@ private SDK buildSDK() {
126227 }
127228
128229 @ CommandLine .Command (name = "decrypt" )
129- void decrypt (@ Option (names = { "-f" , "--file" }, required = true ) Path tdfPath ) throws IOException ,
230+ void decrypt (@ Option (names = { "-f" , "--file" }, required = true ) Path tdfPath ,
231+ @ Option (names = { "--with-assertion-verification-keys" }, defaultValue = Option .NULL_VALUE ) Optional <String > assertionVerification )
232+ throws IOException ,
130233 InvalidAlgorithmParameterException , NoSuchPaddingException , IllegalBlockSizeException ,
131234 BadPaddingException , InvalidKeyException , TDF .FailedToCreateGMAC ,
132235 JOSEException , ParseException , NoSuchAlgorithmException , DecoderException {
133236 var sdk = buildSDK ();
134237 try (var in = FileChannel .open (tdfPath , StandardOpenOption .READ )) {
135238 try (var stdout = new BufferedOutputStream (System .out )) {
136- var reader = new TDF ().loadTDF (in , sdk .getServices ().kas ());
137- reader .readPayload (stdout );
239+ if (assertionVerification .isPresent ()) {
240+ var assertionVerificationInput = assertionVerification .get ();
241+ Gson gson = new Gson ();
242+
243+ AssertionVerificationKeys assertionVerificationKeys ;
244+ try {
245+ assertionVerificationKeys = gson .fromJson (assertionVerificationInput , AssertionVerificationKeys .class );
246+ } catch (JsonSyntaxException e ) {
247+ // try it as a file path
248+ try {
249+ String fileJson = new String (Files .readAllBytes (Paths .get (assertionVerificationInput )));
250+ assertionVerificationKeys = gson .fromJson (fileJson , AssertionVerificationKeys .class );
251+ } catch (JsonSyntaxException e2 ) {
252+ throw new RuntimeException ("Failed to parse assertion verification keys from file" , e2 );
253+ } catch (Exception e3 ) {
254+ throw new RuntimeException ("Could not parse assertion verification keys as json string or path to file" , e3 );
255+ }
256+ }
257+
258+ for (Map .Entry <String , AssertionConfig .AssertionKey > entry : assertionVerificationKeys .keys .entrySet ()){
259+ try {
260+ Object correctedKey = correctKeyType (entry .getValue ().alg , entry .getValue ().key , true );
261+ entry .setValue (new AssertionConfig .AssertionKey (entry .getValue ().alg , correctedKey ));
262+ } catch (Exception e ) {
263+ throw new RuntimeException ("Error with assertion verification key: " + e .getMessage (), e );
264+ }
265+ }
266+ Config .TDFReaderConfig readerConfig = Config .newTDFReaderConfig (
267+ Config .withAssertionVerificationKeys (assertionVerificationKeys ));
268+ var reader = new TDF ().loadTDF (in , sdk .getServices ().kas (), readerConfig );
269+ reader .readPayload (stdout );
270+ } else {
271+ var reader = new TDF ().loadTDF (in , sdk .getServices ().kas ());
272+ reader .readPayload (stdout );
273+ }
138274 }
139275 }
140276 }
0 commit comments