# User Registration and Login using SQLite and Hashlib
# Concept: Salted password hashing and authentication.
import sqlite3
import hashlib
import secrets

def init_db():
    conn = sqlite3.connect("users.db")
    cursor = conn.cursor()
    cursor.execute(
        "CREATE TABLE IF NOT EXISTS users (username TEXT PRIMARY KEY, salt TEXT NOT NULL, password_hash TEXT NOT NULL)"
    )
    conn.commit()
    conn.close()

def hash_password(password, salt=None):
    if not salt:
        salt = secrets.token_hex(16)
    hash_val = hashlib.pbkdf2_hmac("sha256", password.encode(), salt.encode(), 100000)
    return salt, hash_val.hex()

def register_user(username, password):
    conn = sqlite3.connect("users.db")
    cursor = conn.cursor()

    cursor.execute("SELECT * FROM users WHERE username = ?", (username,))
    if cursor.fetchone():
        conn.close()
        return "Username already exists!"

    salt, password_hash = hash_password(password)
    cursor.execute(
        "INSERT INTO users (username, salt, password_hash) VALUES (?, ?, ?)",
        (username, salt, password_hash)
    )
    conn.commit()
    conn.close()
    return "User registered successfully!"

def authenticate_user(username, password):
    conn = sqlite3.connect("users.db")
    cursor = conn.cursor()
    cursor.execute(
        "SELECT salt, password_hash FROM users WHERE username = ?", (username,)
    )
    result = cursor.fetchone()
    conn.close()

    if not result:
        return False

    salt, stored_hash = result
    _, input_hash = hash_password(password, salt)
    return input_hash == stored_hash

if __name__ == "__main__":
    init_db()
    print(register_user("alice", "secret_password"))
    print(register_user("alice", "strong_password"))
    print("Login success?", authenticate_user("alice", "secret_password"))
    print("Login success?", authenticate_user("alice", "wrong_password"))



# Timestamp-Based Replay Attack Detection -- server

import socket
from datetime import datetime, timedelta
from user_auth import authenticate_user, init_db

# Initialize database
init_db()

# Replay protection class
class MessageReceiver:
    def __init__(self):
        self.last_timestamp = None
        self.replay_window = timedelta(seconds=30)

    def verify_timestamp(self, timestamp_str):
        try:
            msg_time = datetime.strptime(timestamp_str, "%Y-%m-%d %H:%M:%S")
            now = datetime.utcnow()

            if self.last_timestamp and msg_time <= self.last_timestamp:
                return False, "Replay attack detected: Old timestamp."

            if abs(now - msg_time) > self.replay_window:
                return False, "Timestamp too old or too far in future."

            self.last_timestamp = msg_time
            return True, "Timestamp is valid."

        except ValueError:
            return False, "Invalid timestamp format."

# Server setup
server = socket.socket()
server.bind(("localhost", 9999))
server.listen(1)

print("Server listening on port 9999...")
conn, addr = server.accept()
print(f"Connected by {addr}")

# Receive login details
username = conn.recv(1024).decode()
password = conn.recv(1024).decode()

if authenticate_user(username, password):
    conn.send(b"AUTH_OK")
    print("Authenticated:", username)

    # Replay protection
    receiver = MessageReceiver()

    while True:
        data = conn.recv(1024).decode()
        if not data:
            break

        try:
            msg, timestamp_str = data.split("||")
        except ValueError:
            conn.send(b"ERROR: Invalid message format")
            continue

        valid, response = receiver.verify_timestamp(timestamp_str)
        print(f"[{username}] {msg} @ {timestamp_str} => {response}")
        conn.send(response.encode())

else:
    conn.send(b"AUTH_FAILED")
    print("Authentication failed.")

conn.close()


# Timestamp-Based Replay Attack Detection -- Client
import socket
from datetime import datetime
import time

client = socket.socket()
client.connect(('localhost', 9999))

# Login
client.send(b"alice")
time.sleep(1)
client.send(b"secret_password")

auth_response = client.recv(1024)

if auth_response != b"AUTH_OK":
    print("Login failed.")
    client.close()
else:
    print("Login successful.")

    # Send messages with timestamps
    while True:
        msg = input("Enter message (or 'exit'): ")

        if msg == "exit":
            break

        timestamp = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")
        full_msg = f"{msg}||{timestamp}"

        client.send(full_msg.encode())
        server_response = client.recv(1024).decode()
        print("Server:", server_response)

        # Simulate replay attempt
        time.sleep(2)
        print("\n[Replaying same message...]")
        client.send(full_msg.encode())
        print("Server:", client.recv(1024).decode())
client.close()




# Encrypted File Transfer over Same Connection -- Server
import socket
from user_auth import authenticate_user, init_db
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

# Constants
AES_KEY = bytes.fromhex("054f0c5c692ffa3e0d0fa7d156966812")
aesgcm = AESGCM(AES_KEY)

# Set up database
init_db()

server = socket.socket()
server.bind(('localhost', 8888))
server.listen(1)

print("File Server listening on port 8888...")
conn, addr = server.accept()
print(f"Connected by {addr}")

# Login
username = conn.recv(1024).decode()
password = conn.recv(1024).decode()

if authenticate_user(username, password):
    conn.send(b"AUTH_OK")
    print("Authenticated:", username)

    # Receive encrypted file
    nonce = conn.recv(12)
    filename = conn.recv(1024).decode()
    file_size = int(conn.recv(1024).decode())

    encrypted_data = b''

    while len(encrypted_data) < file_size:
        chunk = conn.recv(4096)
        if not chunk:
            break
        encrypted_data += chunk

    # Decrypt file data
    try:
        decrypted = aesgcm.decrypt(nonce, encrypted_data, None)

        with open("received_" + filename, 'wb') as f:
            f.write(decrypted)

        conn.send(b"File received and decrypted.")
        print(f"File saved: received_{filename}")

    except Exception as e:
        print("Decryption failed:", e)
        conn.send(b"File decryption failed.")

else:
    conn.send(b"AUTH_FAILED")

conn.close()



# Encrypted File Transfer over Same Connection -- client
import socket
import time
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
import os

# AES Key must match server's
AES_KEY = bytes.fromhex("054f0c5c692ffa3e0d0fa7d156966812")
aesgcm = AESGCM(AES_KEY)

client = socket.socket()
client.connect(('localhost', 8888))

# Login
client.send(b"alice")
time.sleep(1)
client.send(b"secret_password")

if client.recv(1024) != b"AUTH_OK":
    print("Login failed.")
    client.close()
    exit()

print("Logged in. Ready to send file.")

# Select and encrypt file
filename = "sample.txt"

with open(filename, 'rb') as f:
    file_data = f.read()

nonce = os.urandom(12)
encrypted_data = aesgcm.encrypt(nonce, file_data, None)

# Send file
client.send(nonce)
client.send(filename.encode())
client.send(str(len(encrypted_data)).encode())
client.sendall(encrypted_data)

# Receive server response
print("Server:", client.recv(1024).decode())

client.close()





# Generate RSA Public and Private Keys

from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization

# Generate RSA private key
private_key = rsa.generate_private_key(
    public_exponent=65537,
    key_size=2048
)

# Get public key
public_key = private_key.public_key()

# Save private key
with open("client_private.pem", "wb") as f:
    f.write(private_key.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.TraditionalOpenSSL,
        serialization.NoEncryption()
    ))

# Save public key
with open("client_public.pem", "wb") as f:
    f.write(public_key.public_bytes(
        serialization.Encoding.PEM,
        serialization.PublicFormat.SubjectPublicKeyInfo
    ))




# RSA Public-Key Authentication -- Server
import socket
import secrets
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import serialization, hashes

# Load client public key
with open("client_public.pem", "rb") as f:
    client_public_key = serialization.load_pem_public_key(f.read())

# Set up server
server = socket.socket()
server.bind(('localhost', 7777))
server.listen(1)

print("RSA Server running...")

# Accept client connection
conn, addr = server.accept()
print("Connected by", addr)

# Send challenge to client
challenge = secrets.token_bytes(32)
conn.send(challenge)

# Receive signature
signature = conn.recv(512)

# Verify signature
try:
    client_public_key.verify(
        signature,
        challenge,
        padding.PKCS1v15(),
        hashes.SHA256()
    )

    conn.send(b"AUTH_OK")
    print("RSA Authentication Successful")

except Exception as e:
    conn.send(b"AUTH_FAILED")
    print("RSA Authentication Failed:", str(e))

# Close connection
conn.close()
server.close()




# RSA Public-Key Authentication -- Client
import socket
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import serialization, hashes

# Load private key
with open("client_private.pem", "rb") as f:
    private_key = serialization.load_pem_private_key(f.read(), password=None)

# Create socket client
client = socket.socket()
client.connect(('localhost', 7777))

# Receive challenge
challenge = client.recv(1024)

# Sign the challenge
signature = private_key.sign(
    challenge,
    padding.PKCS1v15(),
    hashes.SHA256()
)

# Send signature
client.send(signature)

# Receive response
response = client.recv(1024).decode()
print("Server:", response)

# Close connection
client.close()