Files
zulip/zilencer/management/commands/manage_push_registration_encryption_keys.py
2025-10-17 15:34:34 -07:00

87 lines
3.2 KiB
Python

import configparser
import json
from argparse import ArgumentParser
from typing import Any
from django.conf import settings
from nacl.encoding import Base64Encoder
from nacl.public import PrivateKey
from typing_extensions import override
from zerver.lib.management import ZulipBaseCommand
class Command(ZulipBaseCommand):
help = """
Add or remove a key pair from the `push_registration_encryption_keys` map.
Usage:
./manage.py manage_push_registration_encryption_keys --add
./manage.py manage_push_registration_encryption_keys --remove-key <public-key>
"""
@override
def add_arguments(self, parser: ArgumentParser) -> None:
parser.add_argument(
"--add",
action="store_true",
help="Add a new key pair to the `push_registration_encryption_keys` map.",
)
parser.add_argument(
"--remove-key",
metavar="PUBLIC_KEY",
help="Remove the key pair associated with the given public key from the `push_registration_encryption_keys` map.",
)
@override
def handle(self, *args: Any, **options: Any) -> None:
if not options["add"] and options["remove_key"] is None:
print("Error: Please provide either --add or --remove-key <public-key>.")
return
if settings.DEVELOPMENT:
SECRETS_FILENAME = "zproject/dev-secrets.conf"
else:
SECRETS_FILENAME = "/etc/zulip/zulip-secrets.conf"
config = configparser.ConfigParser()
config.read(SECRETS_FILENAME)
push_registration_encryption_keys: dict[str, str] = json.loads(
config.get("secrets", "push_registration_encryption_keys", fallback="{}")
)
added_key_pair: tuple[str, str] | None = None
if options["add"]:
# Generate a new key-pair and store.
private_key = PrivateKey.generate()
private_key_str = Base64Encoder.encode(bytes(private_key)).decode("utf-8")
public_key_str = Base64Encoder.encode(bytes(private_key.public_key)).decode("utf-8")
push_registration_encryption_keys[public_key_str] = private_key_str
added_key_pair = (public_key_str, private_key_str)
if options["remove_key"] is not None:
# Remove the key-pair for the given public key.
remove_key = options["remove_key"]
if remove_key not in push_registration_encryption_keys:
print("Error: No key pair found for the given public key.")
return
del push_registration_encryption_keys[remove_key]
config.set(
"secrets",
"push_registration_encryption_keys",
json.dumps(push_registration_encryption_keys),
)
with open(SECRETS_FILENAME, "w") as secrets_file:
config.write(secrets_file)
if added_key_pair is not None:
public_key_str, private_key_str = added_key_pair
print("Added a new key pair:")
print(f"- Public key: {public_key_str}")
print(f"- Private key: {private_key_str}")
if options["remove_key"] is not None:
print(f"Removed the key pair for public key: {options['remove_key']}")