From f11905395395e29762438daa431d479a3953f5ff Mon Sep 17 00:00:00 2001
From: Stefan Stammberger <some.fusion@gmail.com>
Date: Tue, 24 Aug 2021 21:23:18 +0200
Subject: [PATCH] fix: main page and creating a user and a wallet

The wallet page will still not renders correctly, but the backend does
create the user his first wallet.
---
 lnbits/core/crud.py          | 20 ++++++----------
 lnbits/core/models.py        | 32 +++++++++++++-------------
 lnbits/core/views/generic.py | 44 ++++++++++++++++++++----------------
 3 files changed, 47 insertions(+), 49 deletions(-)

diff --git a/lnbits/core/crud.py b/lnbits/core/crud.py
index 5378d5933..e981f22c7 100644
--- a/lnbits/core/crud.py
+++ b/lnbits/core/crud.py
@@ -54,20 +54,14 @@ async def get_user(user_id: str, conn: Optional[Connection] = None) -> Optional[
             """,
             (user_id,),
         )
+    else:
+        return None
 
-    return (
-        User(
-            **{
-                **user,
-                **{
-                    "extensions": [e[0] for e in extensions],
-                    "wallets": [Wallet(**w) for w in wallets],
-                },
-            }
-        )
-        if user
-        else None
-    )
+    return User(
+        id = user['id'],
+        email = user['email'],
+        extensions = [e[0] for e in extensions],
+        wallets = [Wallet(**w) for w in wallets])
 
 
 async def update_user_extension(
diff --git a/lnbits/core/models.py b/lnbits/core/models.py
index 8b6ec6ade..c1dc99f91 100644
--- a/lnbits/core/models.py
+++ b/lnbits/core/models.py
@@ -10,22 +10,6 @@ from pydantic import BaseModel
 from lnbits.settings import WALLET
 
 
-class User(BaseModel):
-    id: str
-    email: str
-    extensions: List[str] = []
-    wallets: List["Wallet"] = []
-    password: Optional[str] = None
-
-    @property
-    def wallet_ids(self) -> List[str]:
-        return [wallet.id for wallet in self.wallets]
-
-    def get_wallet(self, wallet_id: str) -> Optional["Wallet"]:
-        w = [wallet for wallet in self.wallets if wallet.id == wallet_id]
-        return w[0] if w else None
-
-
 class Wallet(BaseModel):
     id: str
     name: str
@@ -73,6 +57,22 @@ class Wallet(BaseModel):
         return await get_wallet_payment(self.id, payment_hash)
 
 
+class User(BaseModel):
+    id: str
+    email: Optional[str] = None
+    extensions: List[str] = []
+    wallets: List[Wallet] = []
+    password: Optional[str] = None
+
+    @property
+    def wallet_ids(self) -> List[str]:
+        return [wallet.id for wallet in self.wallets]
+
+    def get_wallet(self, wallet_id: str) -> Optional["Wallet"]:
+        w = [wallet for wallet in self.wallets if wallet.id == wallet_id]
+        return w[0] if w else None
+
+
 class Payment(BaseModel):
     checking_id: str
     pending: bool
diff --git a/lnbits/core/views/generic.py b/lnbits/core/views/generic.py
index fa77030a7..516bcd5c2 100644
--- a/lnbits/core/views/generic.py
+++ b/lnbits/core/views/generic.py
@@ -1,5 +1,8 @@
+from lnbits.core.models import Wallet
 from fastapi.params import Query
 from fastapi.routing import APIRouter
+from fastapi.responses import RedirectResponse
+from fastapi import status
 from lnbits.requestvars import g
 from os import path
 from http import HTTPStatus
@@ -60,11 +63,15 @@ async def extensions(enable: str, disable: str):
     return await templates.TemplateResponse("core/extensions.html", {"request": request, "user": get_user(g.user.id)})
 
 
-@core_html_routes.get("/wallet")
+@core_html_routes.get("/wallet", response_class=HTMLResponse)
 #Not sure how to validate
 @validate_uuids(["usr", "wal"])
-async def wallet(request: Request, usr: Optional[str] = Query(None), 
-                    wal: Optional[str]=Query(None, description=""), nme: Optional[str]=Query(None)):
+async def wallet(request: Request,
+    usr: Optional[str] = Query(None),
+    wal: Optional[str] = Query(None),
+    nme: Optional[str] = Query(None),
+    ):
+                    
     user_id = usr
     wallet_id = wal
     wallet_name = nme
@@ -77,30 +84,27 @@ async def wallet(request: Request, usr: Optional[str] = Query(None),
     # nothing: create everything
 
     if not user_id:
-        user = await get_user((await create_account()).id)
+        usr = await get_user((await create_account()).id)
     else:
-        user = await get_user(user_id)
-        if not user:
-            abort(HTTPStatus.NOT_FOUND, "User does not exist.")
-            return
-
+        usr = await get_user(user_id)
+        if not usr:
+            return g().templates.TemplateResponse("error.html", {"request": request, "err": "User does not exist."})
         if LNBITS_ALLOWED_USERS and user_id not in LNBITS_ALLOWED_USERS:
-            abort(HTTPStatus.UNAUTHORIZED, "User not authorized.")
-
+            return g().templates.TemplateResponse("error.html", {"request": request, "err": "User not authorized."})
     if not wallet_id:
-        if user.wallets and not wallet_name:
-            wallet = user.wallets[0]
+        if usr.wallets and not wallet_name:
+            wal = usr.wallets[0]
         else:
-            wallet = await create_wallet(user_id=user.id, wallet_name=wallet_name)
+            wal = await create_wallet(user_id=usr.id, wallet_name=wallet_name)
 
-        return redirect(url_for("core.wallet", usr=user.id, wal=wallet.id))
+        return RedirectResponse(f"/wallet?usr={usr.id}&wal={wal.id}", status_code=status.HTTP_307_TEMPORARY_REDIRECT)
 
-    wallet = user.get_wallet(wallet_id)
-    if not wallet:
-        abort(HTTPStatus.FORBIDDEN, "Not your wallet.")
+    wal = usr.get_wallet(wallet_id)
+    if not wal:
+        return g().templates.TemplateResponse("error.html", {"request": request, ...})
 
-    return await templates.TemplateResponse(
-        "core/wallet.html", {"request":request,"user":user, "wallet":wallet, "service_fee":service_fee}
+    return g().templates.TemplateResponse(
+        "core/wallet.html", {"request":request,"user":usr, "wallet":wal, "service_fee":service_fee}
     )