pki_ca/ca_core/entity.py

import random
import string

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

    if owner_id is None:
        # Entity is a creator → allow only self
        if requesting_creator_id != entity_id_db:
            raise ValueError("Creator ID does not match entity owner")
    else:
        # Entity is a person/group → allow original creator
        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, group_p, public_key, creator)
        VALUES (%s, FALSE, %s, NULL)
        RETURNING id
        """,
        (name, public_key)
    )
    return cursor.fetchone()["id"]


def enroll_person(cursor, name, public_key, creator_id):
    cursor.execute("SELECT creator FROM entity WHERE id=%s", (creator_id,))
    row = cursor.fetchone()
    if not row or row["creator"] is not None:
        raise ValueError("Provided creator_id does not correspond to a creator")
    cursor.execute(
        """
        INSERT INTO entity (name, group_p, public_key, creator)
        VALUES (%s, FALSE, %s, %s)
        RETURNING id
        """,
        (name, public_key, creator_id)
    )
    return cursor.fetchone()["id"]


def create_group(cursor, name, public_key, creator_id):
    cursor.execute("SELECT creator FROM entity WHERE id=%s", (creator_id,))
    row = cursor.fetchone()
    if not row or row["creator"] is not None:
        raise ValueError("Provided creator_id does not correspond to a creator")
    cursor.execute(
        """
        INSERT INTO entity (name, group_p, public_key, creator)
        VALUES (%s, TRUE, %s, %s)
        RETURNING id
        """,
        (name, public_key, creator_id)
    )
    return cursor.fetchone()["id"]


def create_alias(cursor, person_id):
    cursor.execute("SELECT id, group_p, public_key FROM entity WHERE id=%s", (person_id,))
    row = cursor.fetchone()
    if not row:
        raise ValueError("Person not found")
    if row["group_p"]:
        raise ValueError("Groups cannot create aliases")
    random_name = ''.join(random.choices(string.ascii_letters + string.digits, k=8))
    cursor.execute(
        """
        INSERT INTO entity (name, group_p, public_key, creator)
        VALUES (%s, FALSE, %s, %s)
        RETURNING id
        """,
        (random_name, row["public_key"], person_id)
    )
    return cursor.fetchone()["id"]


# ------------------------
# Deletion
# ------------------------
def delete_person(cursor, person_id, requesting_creator_id):
    _verify_ownership(cursor, person_id, requesting_creator_id)
    cursor.execute("DELETE FROM entity WHERE id=%s AND group_p=FALSE", (person_id,))
    if cursor.rowcount == 0:
        raise ValueError("Person not found or already deleted")


def delete_group(cursor, group_id, requesting_creator_id):
    _verify_ownership(cursor, group_id, requesting_creator_id)
    cursor.execute("DELETE FROM entity WHERE id=%s AND group_p=TRUE", (group_id,))
    if cursor.rowcount == 0:
        raise ValueError("Group not found or already deleted")


def delete_creator(cursor, creator_id, requesting_creator_id):
    _verify_ownership(cursor, creator_id, requesting_creator_id)
    # check if creator has created other entities
    cursor.execute("SELECT COUNT(*) as cnt FROM entity WHERE creator=%s", (creator_id,))
    if cursor.fetchone()["cnt"] > 0:
        raise ValueError("Creator cannot be deleted because it has created other entities")
    cursor.execute("DELETE FROM entity WHERE id=%s AND creator IS NULL", (creator_id,))
    if cursor.rowcount == 0:
        raise ValueError("Creator not found or already deleted")


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


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


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


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


# ------------------------
# Setters
# ------------------------
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):
    # only public_key for current schema
    set_entity_public_key(cursor, entity_id, public_key, requesting_creator_id)