pki_ca/ca_core/entity.py

import random
import string

# ------------------------
# Helper for ownership checks
# ------------------------
def _verify_ownership(cursor, entity_id, requesting_creator_id):
    cursor.execute(
        "SELECT id, creator, type, status FROM entity WHERE id=%s", (entity_id,)
    )
    row = cursor.fetchone()
    if not row or row["status"] != "active":
        raise ValueError("Entity not found or inactive")
    owner_id = row["creator"]
    entity_type = row["type"]
    entity_id_db = row["id"]

    if entity_type == "creator":
        if requesting_creator_id != entity_id_db:
            raise ValueError("Creator ID does not match entity owner")
    else:
        if requesting_creator_id != owner_id:
            raise ValueError("Creator ID does not match entity owner")


# ------------------------
# Insertions
# ------------------------
def insert_creator(cursor, name, public_key):
    cursor.execute(
        """
        INSERT INTO entity (name, type, public_key, creator, status)
        VALUES (%s, 'creator', %s, NULL, 'active')
        RETURNING id
        """,
        (name, public_key)
    )
    return cursor.fetchone()["id"]


def enroll_person(cursor, name, public_key, creator_id):
    cursor.execute(
        "SELECT type, status FROM entity WHERE id=%s", (creator_id,)
    )
    row = cursor.fetchone()
    if not row or row["type"] != "creator" or row["status"] != "active":
        raise ValueError("Provided creator_id does not correspond to a valid active creator")

    cursor.execute(
        """
        INSERT INTO entity (name, type, public_key, creator, status)
        VALUES (%s, 'person', %s, %s, 'active')
        RETURNING id
        """,
        (name, public_key, creator_id)
    )
    return cursor.fetchone()["id"]


def create_group(cursor, name, public_key, creator_id):
    cursor.execute(
        "SELECT type, status FROM entity WHERE id=%s", (creator_id,)
    )
    row = cursor.fetchone()
    if not row or row["type"] != "creator" or row["status"] != "active":
        raise ValueError("Provided creator_id does not correspond to a valid active creator")

    cursor.execute(
        """
        INSERT INTO entity (name, type, public_key, creator, status)
        VALUES (%s, 'group', %s, %s, 'active')
        RETURNING id
        """,
        (name, public_key, creator_id)
    )
    return cursor.fetchone()["id"]


def create_alias(cursor, person_id):
    cursor.execute(
        "SELECT id, type, public_key, status FROM entity WHERE id=%s", (person_id,)
    )
    row = cursor.fetchone()
    if not row or row["status"] != "active":
        raise ValueError("Person not found or inactive")
    if row["type"] != "person":
        raise ValueError("Only persons can create aliases")

    random_name = "".join(random.choices(string.ascii_letters + string.digits, k=8))

    cursor.execute(
        """
        INSERT INTO entity (name, type, public_key, creator, status)
        VALUES (%s, 'person', %s, %s, 'active')
        RETURNING id
        """,
        (random_name, row["public_key"], person_id)
    )
    return cursor.fetchone()["id"]


# ------------------------
# Soft-delete / revocation
# ------------------------
def revoke_entity(cursor, entity_id, requesting_creator_id):
    _verify_ownership(cursor, entity_id, requesting_creator_id)
    cursor.execute(
        "UPDATE entity SET status=%s WHERE id=%s", ("revoked", entity_id)
    )


# ------------------------
# Getters / Setters
# ------------------------
def get_entity(cursor, entity_id):
    cursor.execute(
        "SELECT * FROM entity WHERE id=%s AND status='active'", (entity_id,)
    )
    row = cursor.fetchone()
    if not row:
        raise ValueError("Entity not found or inactive")
    return row


def get_entity_id(cursor, name):
    cursor.execute(
        "SELECT id FROM entity WHERE name=%s AND status='active'", (name,)
    )
    row = cursor.fetchone()
    if not row:
        raise ValueError("Entity not found or inactive")
    return row["id"]


def get_entity_public_key(cursor, entity_id):
    cursor.execute(
        "SELECT public_key FROM entity WHERE id=%s AND status='active'", (entity_id,)
    )
    row = cursor.fetchone()
    if not row:
        raise ValueError("Entity not found or inactive")
    return row["public_key"]


def get_entity_name(cursor, entity_id):
    cursor.execute(
        "SELECT name FROM entity WHERE id=%s AND status='active'", (entity_id,)
    )
    row = cursor.fetchone()
    if not row:
        raise ValueError("Entity not found or inactive")
    return row["name"]


def set_entity_name(cursor, entity_id, new_name, requesting_creator_id):
    _verify_ownership(cursor, entity_id, requesting_creator_id)
    cursor.execute("UPDATE entity SET name=%s WHERE id=%s", (new_name, entity_id))


def set_entity_public_key(cursor, entity_id, public_key, requesting_creator_id):
    _verify_ownership(cursor, entity_id, requesting_creator_id)
    cursor.execute(
        "UPDATE entity SET public_key=%s WHERE id=%s", (public_key, entity_id)
    )


def set_entity_keys(cursor, entity_id, public_key, requesting_creator_id):
    set_entity_public_key(cursor, entity_id, public_key, requesting_creator_id)