xref: /aosp_15_r20/external/tink/python/examples/aead/aead.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 aead-example]
15"""A command-line utility for encrypting small files with AEAD.
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, ['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.')
37flags.DEFINE_string('associated_data', None,
38                    'Optional associated data used for the encryption.')
39
40
41def main(argv):
42  del argv  # Unused.
43
44  associated_data = b'' if not FLAGS.associated_data else bytes(
45      FLAGS.associated_data, 'utf-8')
46
47  # Initialise Tink
48  aead.register()
49
50  # Read the keyset into a keyset_handle
51  with open(FLAGS.keyset_path, 'rt') as keyset_file:
52    try:
53      text = keyset_file.read()
54      keyset_handle = cleartext_keyset_handle.read(tink.JsonKeysetReader(text))
55    except tink.TinkError as e:
56      logging.exception('Error reading key: %s', e)
57      return 1
58
59  # Get the primitive
60  try:
61    cipher = keyset_handle.primitive(aead.Aead)
62  except tink.TinkError as e:
63    logging.error('Error creating primitive: %s', e)
64    return 1
65
66  with open(FLAGS.input_path, 'rb') as input_file:
67    input_data = input_file.read()
68    if FLAGS.mode == 'decrypt':
69      output_data = cipher.decrypt(input_data, associated_data)
70    elif FLAGS.mode == 'encrypt':
71      output_data = cipher.encrypt(input_data, associated_data)
72    else:
73      logging.error(
74          'Error mode not supported. Please choose "encrypt" or "decrypt".')
75      return 1
76
77    with open(FLAGS.output_path, 'wb') as output_file:
78      output_file.write(output_data)
79
80
81if __name__ == '__main__':
82  flags.mark_flags_as_required([
83      'mode', 'keyset_path', 'input_path', 'output_path'])
84  app.run(main)
85# [END aead-example]
86