xref: /aosp_15_r20/external/tink/python/examples/cleartext_keyset/cleartext_keyset.py (revision e7b1675dde1b92d52ec075b0a92829627f2c52a5) 
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# [START cleartext-keyset-example]
15"""A command-line utility for generating, storing and using cleartext AES128_GCM keysets.
16
17It loads cleartext keys from disk - this is not recommended!
18"""
19
20from absl import app
21from absl import flags
22from absl import logging
23
24import tink
25from tink import aead
26from tink import cleartext_keyset_handle
27
28
29FLAGS = flags.FLAGS
30
31flags.DEFINE_enum('mode', None, ['generate', 'encrypt', 'decrypt'],
32                  'The operation to perform.')
33flags.DEFINE_string('keyset_path', None,
34                    'Path to the keyset used for encryption.')
35flags.DEFINE_string('input_path', None, 'Path to the input file.')
36flags.DEFINE_string('output_path', None, 'Path to the output file.')
37
38
39def main(argv):
40  del argv  # Unused.
41
42  # Initialise Tink
43  aead.register()
44
45  if FLAGS.mode == 'generate':
46    # [START generate-a-new-keyset]
47    # Generate a new keyset
48    try:
49      key_template = aead.aead_key_templates.AES128_GCM
50      keyset_handle = tink.new_keyset_handle(key_template)
51    except tink.TinkError as e:
52      logging.exception('Error creating primitive: %s', e)
53      return 1
54    # [END generate-a-new-keyset]
55
56    # [START store-a-cleartext-keyset]
57    with open(FLAGS.keyset_path, 'wt') as keyset_file:
58      try:
59        cleartext_keyset_handle.write(
60            tink.JsonKeysetWriter(keyset_file), keyset_handle)
61      except tink.TinkError as e:
62        logging.exception('Error writing key: %s', e)
63        return 1
64    return 0
65    # [END store-a-cleartext-keyset]
66
67  # Use the input keyset to encrypt/decrypt data
68
69  # Read the keyset into a keyset_handle
70  with open(FLAGS.keyset_path, 'rt') as keyset_file:
71    try:
72      text = keyset_file.read()
73      keyset_handle = cleartext_keyset_handle.read(tink.JsonKeysetReader(text))
74    except tink.TinkError as e:
75      logging.exception('Error reading key: %s', e)
76      return 1
77
78  # Get the primitive
79  try:
80    cipher = keyset_handle.primitive(aead.Aead)
81  except tink.TinkError as e:
82    logging.error('Error creating primitive: %s', e)
83    return 1
84
85  with open(FLAGS.input_path, 'rb') as input_file:
86    input_data = input_file.read()
87    if FLAGS.mode == 'decrypt':
88      output_data = cipher.decrypt(input_data, b'envelope_example')
89    elif FLAGS.mode == 'encrypt':
90      output_data = cipher.encrypt(input_data, b'envelope_example')
91    else:
92      logging.error(
93          'Error mode not supported. Please choose "encrypt" or "decrypt".')
94      return 1
95
96    with open(FLAGS.output_path, 'wb') as output_file:
97      output_file.write(output_data)
98
99if __name__ == '__main__':
100  flags.mark_flags_as_required(['mode', 'keyset_path'])
101  app.run(main)
102# [END cleartext-keyset-example]
103