Skip to content

Python library that implements the PHT security protocol for train containers and images

License

Notifications You must be signed in to change notification settings

PHT-Medic/train-container-library

Repository files navigation

Documentation Status CodeQL main-ci codecov PyPI - Python Version PyPI - Downloads Code style: black

🚆 Train Container Library

Python library for validating and interacting with pht-train images/containers.

Installation

pip install pht-train-container-library

Setup development environment

Make sure you have poetry and pre-commit installed.

Install the dependencies and pre-commit hooks:

poetry install --with dev
poetry run pre-commit install

Run tests

poetry run pytest

Linting and formatting

These commands are also run as pre-commit hooks.

Linting with ruff:

poetry run ruff . --fix

Formatting with black:

poetry run black .

Security Protocol

The pht security protocol adapted from docs/Secure_PHT_latest__official.pdf performs two main tasks:

  1. Before executing a train-image on the local machine, unless the station is the first station on the route, the previous results need to be decrypted and the content of the image needs to be validated based on the configuration of the individual train -> pre-run.
  2. After executing the train the updated results need to be encrypted and the train configuration needs to be updated to reflect the current state ->post-run.

Train image structure

To ensure the protocol is working correctly train docker images are required to keep the following structure:

  • /opt/train_config.json: Stores the configuration file of the train.
  • /opt/pht_train/: Stores all the files containing code or other things required for the train algorithm to run. The contents of this directory can never change and is validated by the pre-run step.
  • /opt/pht_results/: Stores the results of the train. Which will be decrypted in the pre-run step and encrypted in the post-run step.

No files in the image outside the /opt/pht_results/ directory should change during the execution of the algorithm.

Usage - Python Script

To use the protocol in your own python application, after installing the library with pip install pht-train-container-library an instance of the protocol can be to validate docker images as follows:

from train_lib.security.protocol import SecurityProtocol
from train_lib.docker_util.docker_ops import extract_train_config

image_name = '<image-repo>:<image-tag>'
station_id = '<station-id>'

# Get the train configuration from the image
config = extract_train_config(image_name)
# Initialize the protocol with the extracted config and station_id
protocol = SecurityProtocol(station_id=station_id, config=config)

# execute one of the protocol steps
protocol.pre_run_protocol(image_name, private_key_path='<path-to-private-key>')
# protocol.post_run_protocol(image_name, private_key_path='<path-to-private-key>')

Usage - Container

A containerized version of the protocol is also available it can be used with the following command:

docker run -e STATION_ID=<station_id> -e PRIVATE_KEY_PATH=/opt/private_key.pem -v /var/run/docker.sock:/var/run/docker.sock -v <path_to_your_key>:/opt/private_key.pem ghcr.io/pht-medic/protocol <pre-run/post-run> <image-repo>:<image-tag>

STATION_ID and PRIVATE_KEY_PATH are required to be set in the environment variables. As well as passing the docker socket /var/run/docker.sock to the container as a volume to enable docker-in-docker functionality.

Pre-run protocol

The pre-run protocol consists of the following steps

  1. The hash of the immutable files (train definition) is verified making sure that the executable files did not change during the the train definition.
  2. The digital signature is verified ensuring the correctness of the results at each stop of the train.
  3. The symmetric key is decrypted using the provided station private key
  4. The mutable files in /opt/pht_results are decrypted using the symmetric key obtained in the previous step
  5. The decrypted files are hashed and the hash is compared to the one stored in the train configuration file.

Once these steps have been completed the image is ready to be executed.

Post-run protocol

  1. Calculate the hash of the newly generated results
  2. Sign the hash of the results using the provided PRIVATE_KEY_PATH
  3. Update the the train signature using the session id that is randomly generated at each execution step
  4. Encrypt the resulting files using a newly generated symmetric key
  5. Encrypt the generated symmetric key with the public keys of the train participants
  6. Update the train configuration file

With the completion of these steps the train is ready to be pushed into the registry for further processing

Tests

Run the tests to validate the security protocol is working as intended. From this projects root directory run pytest train_lib