openvpn-2fa-otp-freeradius-ldap

1
import radiusd
2
from sqlalchemy import select, create_engine, URL, String
3
from sqlalchemy.orm import DeclarativeBase, Session
4
from typing import Optional, TypedDict
5
from sqlalchemy.orm import Mapped
6
from sqlalchemy.orm import mapped_column
7
from config import DBConfig
8
from util import log_msg
9

10
class Base(DeclarativeBase):
11
    pass
12

13

14
class User(Base):
15
    __tablename__ = "users"
16

17
    id: Mapped[int] = mapped_column(primary_key=True)
18
    login: Mapped[str] = mapped_column(String[200])
19
    otp_secret: Mapped[str] = mapped_column(String[50])
20

21

22
class DBConnection:
23
    def __init__(self, config: DBConfig):
24
        self.config = config
25
        self.conn = self.connect()
26

27
    def connect(self):
28
        url_object = URL.create(
29
            "postgresql+psycopg2",
30
            username=self.config["user"],
31
            password=self.config["password"],
32
            host=self.config["host"],
33
            port=self.config["port"],
34
            database=self.config["db_name"],
35
        )
36
        return create_engine(url_object, echo=False)
37

38
    def get_otp_secret(self, login: str) -> Optional[str]:
39
        # select user
40
        try:
41
            with Session(self.conn) as session:
42
                stmt = select(User.otp_secret).filter_by(login=login)
43
                otp_secret = session.execute(stmt).first()
44
                if otp_secret is None:
45
                    return None
46
                return otp_secret[0]
47
        except Exception as e:
48
            log_msg(radiusd.L_ERR,
49
                           "{}".format(e))
50
            return None
51