from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
from typing import List, Optional
from sqlmodel import SQLModel, Field, Relationship, VARCHAR, Column
from pydantic import EmailStr


class UserGroup(SQLModel, table=True):
    id: int = Field(default=None, primary_key=True)
    name: str

    users: List["User"] = Relationship(back_populates="group")


class User(SQLModel, table=True):
    id: int = Field(default=None, primary_key=True)
    username: str = Field(sa_column=Column('username', VARCHAR, unique=True))
    email: EmailStr
    password: str
    first_name: str
    last_name: str
    public_key: str
    private_key: str
    profile_pic: Optional[str]
    is_staff: bool = Field(default=False)

    group_id: Optional[int] = Field(foreign_key="usergroup.id")
    group: Optional[UserGroup] = Relationship(back_populates="users")


class RemoteUser(SQLModel, table=True):
    id: int = Field(default=None, primary_key=True)
    remote_id: str = Field(sa_column=Column('username', VARCHAR, unique=True))
    inbox: str
    public_key: Optional[str]


def generate_new_pk() -> RSAPrivateKey:
    return rsa.generate_private_key(public_exponent=65537, key_size=4096)


def load_pk(pkstr: str) -> RSAPrivateKey:
    loaded_pk = serialization.load_pem_private_key(pkstr, password=None)


def create_user(username: str, password: str, email: str, first_name: str, last_name: str, profile_pic: Optional[str]) -> User:
    user = User(
        username=username,
        password=password,
        email=email,
        first_name=first_name,
        last_name=last_name,
        profile_pic=profile_pic,
    )

    pk = generate_new_pk()
    pub_key = pk.public_key()

    user.public_key = pub_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    ).decode('utf-8')

    user.private_key = pk.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption()
    ).decode('utf-8')

    return user
cryptography.hazmat.primitives.serialization.load_pem_private_key cryptography.hazmat.primitives.serialization.NoEncryption sqlmodel.Relationship sqlmodel.Field sqlmodel.Column cryptography.hazmat.primitives.asymmetric.rsa.generate_private_key
[((374, 411), 'sqlmodel.Field', 'Field', ([], {'default': 'None', 'primary_key': '(True)'}), '(default=None, primary_key=True)\n', (379, 411), False, 'from sqlmodel import SQLModel, Field, Relationship, VARCHAR, Column\n'), ((453, 489), 'sqlmodel.Relationship', 'Relationship', ([], {'back_populates': '"""group"""'}), "(back_populates='group')\n", (465, 489), False, 'from sqlmodel import SQLModel, Field, Relationship, VARCHAR, Column\n'), ((540, 577), 'sqlmodel.Field', 'Field', ([], {'default': 'None', 'primary_key': '(True)'}), '(default=None, primary_key=True)\n', (545, 577), False, 'from sqlmodel import SQLModel, Field, Relationship, VARCHAR, Column\n'), ((826, 846), 'sqlmodel.Field', 'Field', ([], {'default': '(False)'}), '(default=False)\n', (831, 846), False, 'from sqlmodel import SQLModel, Field, Relationship, VARCHAR, Column\n'), ((878, 911), 'sqlmodel.Field', 'Field', ([], {'foreign_key': '"""usergroup.id"""'}), "(foreign_key='usergroup.id')\n", (883, 911), False, 'from sqlmodel import SQLModel, Field, Relationship, VARCHAR, Column\n'), ((945, 981), 'sqlmodel.Relationship', 'Relationship', ([], {'back_populates': '"""users"""'}), "(back_populates='users')\n", (957, 981), False, 'from sqlmodel import SQLModel, Field, Relationship, VARCHAR, Column\n'), ((1038, 1075), 'sqlmodel.Field', 'Field', ([], {'default': 'None', 'primary_key': '(True)'}), '(default=None, primary_key=True)\n', (1043, 1075), False, 'from sqlmodel import SQLModel, Field, Relationship, VARCHAR, Column\n'), ((1253, 1315), 'cryptography.hazmat.primitives.asymmetric.rsa.generate_private_key', 'rsa.generate_private_key', ([], {'public_exponent': '(65537)', 'key_size': '(4096)'}), '(public_exponent=65537, key_size=4096)\n', (1277, 1315), False, 'from cryptography.hazmat.primitives.asymmetric import rsa\n'), ((1376, 1432), 'cryptography.hazmat.primitives.serialization.load_pem_private_key', 'serialization.load_pem_private_key', (['pkstr'], {'password': 'None'}), '(pkstr, password=None)\n', (1410, 1432), False, 'from cryptography.hazmat.primitives import serialization\n'), ((614, 654), 'sqlmodel.Column', 'Column', (['"""username"""', 'VARCHAR'], {'unique': '(True)'}), "('username', VARCHAR, unique=True)\n", (620, 654), False, 'from sqlmodel import SQLModel, Field, Relationship, VARCHAR, Column\n'), ((1113, 1153), 'sqlmodel.Column', 'Column', (['"""username"""', 'VARCHAR'], {'unique': '(True)'}), "('username', VARCHAR, unique=True)\n", (1119, 1153), False, 'from sqlmodel import SQLModel, Field, Relationship, VARCHAR, Column\n'), ((2154, 2182), 'cryptography.hazmat.primitives.serialization.NoEncryption', 'serialization.NoEncryption', ([], {}), '()\n', (2180, 2182), False, 'from cryptography.hazmat.primitives import serialization\n')]