feat: add SRP authentication, improve security
- Replace RSA key exchange with SRP (Secure Remote Password) - Password never transmitted over network - Add unit tests for endpoints - Fix datetime.UTC compatibility for Python < 3.11 - Fix logger.exception usage - Update README with new auth flow diagram
This commit is contained in:
mirai
+55
-35
@@ -1,14 +1,17 @@
|
||||
import asyncio
|
||||
import json
|
||||
import base64
|
||||
from typing import Optional
|
||||
|
||||
import rsa
|
||||
import srp
|
||||
import requests
|
||||
from cryptography.fernet import Fernet
|
||||
import websockets
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
|
||||
srp.rfc5054_enable()
|
||||
|
||||
|
||||
class Client:
|
||||
def __init__(
|
||||
@@ -17,11 +20,8 @@ class Client:
|
||||
self.server = server
|
||||
self.port = port
|
||||
self.username = username
|
||||
self.password = password or ""
|
||||
self.password = (password or "").encode()
|
||||
self.user_id: Optional[str] = None
|
||||
|
||||
self.public_key: Optional[rsa.PublicKey] = None
|
||||
self.private_key: Optional[rsa.PrivateKey] = None
|
||||
self.fernet: Optional[Fernet] = None
|
||||
|
||||
self.console = Console()
|
||||
@@ -47,34 +47,55 @@ class Client:
|
||||
def info(self, message: str) -> None:
|
||||
self.console.print(f"[cyan]• {message}[/]")
|
||||
|
||||
def generate_keys(self) -> None:
|
||||
with self.console.status(
|
||||
"[cyan]Generating RSA keys (2048 bit)...[/]", spinner="dots"
|
||||
):
|
||||
self.public_key, self.private_key = rsa.newkeys(2048)
|
||||
self.success("RSA keys generated")
|
||||
def srp_authenticate(self) -> None:
|
||||
"""SRP authentication flow"""
|
||||
with self.console.status("[cyan]Starting SRP handshake...[/]", spinner="dots"):
|
||||
|
||||
def exchange_keys(self) -> None:
|
||||
with self.console.status(
|
||||
"[cyan]Exchanging keys with server...[/]", spinner="dots"
|
||||
):
|
||||
pubkey_bytes = self.public_key.save_pkcs1()
|
||||
response = requests.post(
|
||||
f"{self.base_url}/get_key",
|
||||
files={"pubkey": ("key.pem", pubkey_bytes)},
|
||||
data={"username": self.username, "password": self.password},
|
||||
usr = srp.User(b"chat", self.password, hash_alg=srp.SHA256)
|
||||
_, A = usr.start_authentication()
|
||||
|
||||
resp = requests.post(
|
||||
f"{self.base_url}/srp/init",
|
||||
json={
|
||||
"username": self.username,
|
||||
"A": base64.b64encode(A).decode(),
|
||||
},
|
||||
timeout=30,
|
||||
)
|
||||
response.raise_for_status()
|
||||
resp.raise_for_status()
|
||||
init_data = resp.json()
|
||||
|
||||
self.user_id = response.headers.get("X-User-Id")
|
||||
encrypted_key = response.content
|
||||
symmetric_key = rsa.decrypt(encrypted_key, self.private_key)
|
||||
self.fernet = Fernet(symmetric_key)
|
||||
self.user_id = init_data["user_id"]
|
||||
B = base64.b64decode(init_data["B"])
|
||||
salt = base64.b64decode(init_data["salt"])
|
||||
|
||||
self.success(f"Key exchange complete (session: {self.user_id[:8]}...)")
|
||||
self.public_key = None
|
||||
self.private_key = None
|
||||
M = usr.process_challenge(salt, B)
|
||||
|
||||
if M is None:
|
||||
raise ValueError("SRP challenge processing failed")
|
||||
|
||||
resp = requests.post(
|
||||
f"{self.base_url}/srp/verify",
|
||||
json={
|
||||
"user_id": self.user_id,
|
||||
"username": self.username,
|
||||
"M": base64.b64encode(M).decode(),
|
||||
},
|
||||
timeout=30,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
verify_data = resp.json()
|
||||
|
||||
H_AMK = base64.b64decode(verify_data["H_AMK"])
|
||||
usr.verify_session(H_AMK)
|
||||
|
||||
if not usr.authenticated():
|
||||
raise ValueError("Server authentication failed")
|
||||
|
||||
session_key = base64.b64decode(verify_data["session_key"])
|
||||
self.fernet = Fernet(session_key)
|
||||
|
||||
self.success(f"SRP authenticated (session: {self.user_id[:8]}...)")
|
||||
|
||||
def render_messages(self) -> None:
|
||||
self.console.clear()
|
||||
@@ -147,13 +168,10 @@ class Client:
|
||||
self.console.print()
|
||||
|
||||
try:
|
||||
self.generate_keys()
|
||||
self.exchange_keys()
|
||||
self.srp_authenticate()
|
||||
|
||||
self.info("Connecting to chat...")
|
||||
url = (
|
||||
f"{self.ws_url}/ws/chat?user_id={self.user_id}&password={self.password}"
|
||||
)
|
||||
url = f"{self.ws_url}/ws/chat?user_id={self.user_id}"
|
||||
|
||||
async with websockets.connect(url) as ws:
|
||||
self.success("Connected to chat server")
|
||||
@@ -175,10 +193,12 @@ class Client:
|
||||
self.error(f"Cannot connect to {self.base_url}")
|
||||
except requests.exceptions.HTTPError as e:
|
||||
self.error(f"Server error: {e.response.status_code} - {e.response.text}")
|
||||
except Exception as e:
|
||||
except ValueError as e:
|
||||
self.error(f"Authentication failed: {e}")
|
||||
except Exception:
|
||||
import traceback
|
||||
|
||||
self.error(f"Error: {e}")
|
||||
self.error("Error occurred")
|
||||
traceback.print_exc()
|
||||
|
||||
def run(self) -> None:
|
||||
|
||||
@@ -1 +1 @@
|
||||
MESSAGES_TO_SHOW = 5
|
||||
MESSAGES_TO_SHOW = 5
|
||||
|
||||
@@ -1,28 +0,0 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class CryptoService(ABC):
|
||||
|
||||
@abstractmethod
|
||||
def _encrypt(self, message: str) -> str:
|
||||
raise NotImplementedError("Need to implement encrypt method")
|
||||
|
||||
@abstractmethod
|
||||
def _decrypt(self, message: str) -> str:
|
||||
raise NotImplementedError("Need to implement decrypt method")
|
||||
|
||||
@abstractmethod
|
||||
def _request_key(self, url: str, username: str, password: str | None = None):
|
||||
raise NotImplementedError("Need to implement request key method")
|
||||
|
||||
@abstractmethod
|
||||
def _generate_keys(self):
|
||||
raise NotImplementedError("Need to implement generate keys method")
|
||||
|
||||
@abstractmethod
|
||||
def _get_generated_keys(self) -> tuple:
|
||||
raise NotImplementedError("Need to implement get generated keys method")
|
||||
|
||||
@abstractmethod
|
||||
def _remove_keys(self):
|
||||
raise NotImplementedError("Need to implement remove keys method")
|
||||
@@ -1,33 +0,0 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class ClientRenderer(ABC):
|
||||
|
||||
username: str
|
||||
|
||||
@abstractmethod
|
||||
def _decrypt(self, message: str) -> str:
|
||||
"""Decrypt an encrypted message (provided by crypto mixin)."""
|
||||
raise NotImplementedError("Need to implement _decrypt")
|
||||
|
||||
@abstractmethod
|
||||
def print_message(self, message: str) -> str:
|
||||
raise NotImplementedError("Need to implement print_message")
|
||||
|
||||
@abstractmethod
|
||||
def clear_console(self) -> None:
|
||||
"""Clear the client console (platform-specific)."""
|
||||
raise NotImplementedError("Need to implement clear_console")
|
||||
|
||||
@abstractmethod
|
||||
def print_ip(self, ip: str) -> str:
|
||||
raise NotImplementedError("Need to implement print_ip")
|
||||
|
||||
@abstractmethod
|
||||
def print_username(self, username: str) -> str:
|
||||
raise NotImplementedError("Need to implement print_username")
|
||||
|
||||
@abstractmethod
|
||||
def print_chat(self, response) -> None:
|
||||
"""Render chat payload (response is expected to be a mapping with 'messages' and 'users_in_chat')."""
|
||||
raise NotImplementedError("Need to implement print_chat")
|
||||
@@ -1,43 +0,0 @@
|
||||
import rsa
|
||||
import requests
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
from cmd_chat.client.core.abs.abs_crypto import CryptoService
|
||||
|
||||
|
||||
class RSAService(CryptoService):
|
||||
def __init__(self):
|
||||
self.public_key = None
|
||||
self.private_key = None
|
||||
self.symmetric_key = None
|
||||
self.fernet = None
|
||||
self._generate_keys()
|
||||
|
||||
def _encrypt(self, message: str) -> str:
|
||||
return self.fernet.encrypt(message.encode()).decode("utf-8")
|
||||
|
||||
def _decrypt(self, message: str) -> str:
|
||||
return self.fernet.decrypt(message.encode()).decode("utf-8")
|
||||
|
||||
def _request_key(self, url: str, username: str, password: str | None = None):
|
||||
pubkey_bytes = self.public_key.save_pkcs1()
|
||||
r = requests.post(
|
||||
url,
|
||||
files={"pubkey": ("public.pem", pubkey_bytes)},
|
||||
data={"username": username, "password": password or ""},
|
||||
stream=True,
|
||||
)
|
||||
r.raise_for_status()
|
||||
message = r.content
|
||||
self.symmetric_key = rsa.decrypt(message, self.private_key)
|
||||
self.fernet = Fernet(self.symmetric_key)
|
||||
|
||||
def _generate_keys(self):
|
||||
self.public_key, self.private_key = rsa.newkeys(2048)
|
||||
|
||||
def _get_generated_keys(self):
|
||||
return self.private_key, self.public_key
|
||||
|
||||
def _remove_keys(self):
|
||||
self.public_key = None
|
||||
self.private_key = None
|
||||
@@ -1,72 +0,0 @@
|
||||
import os
|
||||
import platform
|
||||
|
||||
from rich.text import Text
|
||||
from rich.style import Style
|
||||
from rich.console import Console
|
||||
|
||||
from rich.table import Table
|
||||
from cmd_chat.client.core.abs.abs_renderer import ClientRenderer
|
||||
from cmd_chat.client.config import MESSAGES_TO_SHOW
|
||||
|
||||
|
||||
console = Console(width=75)
|
||||
|
||||
|
||||
class RichClientRenderer(ClientRenderer):
|
||||
|
||||
def __get_os(self) -> str:
|
||||
"""checking what kind of platform you need"""
|
||||
if "Linux" in str(platform.platform()):
|
||||
return "Linux"
|
||||
return "Windows"
|
||||
|
||||
def print_message(self, message: str) -> Text:
|
||||
"""generating string with message in required format"""
|
||||
# split only on the first ':' so message bodies containing ':' are preserved
|
||||
parts = message.split(":", 1)
|
||||
if parts[0] == self.username:
|
||||
return (
|
||||
Text(text=parts[0], style="bold")
|
||||
+ Text(text=": ", style="bold")
|
||||
+ Text(text=parts[1], style="underline")
|
||||
)
|
||||
return (
|
||||
Text(text=parts[0], style="bold")
|
||||
+ Text(text=": ", style="bold")
|
||||
+ Text(text=parts[1], style="underline")
|
||||
)
|
||||
|
||||
def clear_console(self):
|
||||
# For windows clear command its cls
|
||||
# For linux clear command its clear
|
||||
if self.__get_os() == "Linux":
|
||||
os.system("clear")
|
||||
else:
|
||||
os.system("cls")
|
||||
|
||||
def print_ip(self, ip: str) -> str:
|
||||
return ip
|
||||
|
||||
def print_username(self, username: str) -> str:
|
||||
return username
|
||||
|
||||
def print_chat(self, response) -> None:
|
||||
self.clear_console()
|
||||
for i, msg in enumerate(response["messages"][-MESSAGES_TO_SHOW:]):
|
||||
actual_message = self._decrypt(msg)
|
||||
if i == 0:
|
||||
console.print("Users in chat:", justify="left")
|
||||
table = Table(show_header=True, header_style="bold magenta")
|
||||
table.add_column("IP", style="dim", width=12)
|
||||
table.add_column("USERNAME")
|
||||
for user in response["users_in_chat"]:
|
||||
table.add_row(
|
||||
self.print_ip(user.split(",")[0]),
|
||||
self.print_username(user.split(",")[1]),
|
||||
)
|
||||
console.print(table)
|
||||
console.print("Write 'q' to quit from chat", justify="left")
|
||||
console.print(f"\n{self.print_message(actual_message)}")
|
||||
else:
|
||||
console.print(f"{self.print_message(actual_message)}")
|
||||
Reference in New Issue
Block a user