bharathmunakala/Role_Base_Access_Control

1import os
2from datetime import datetime
3from typing import Optional, Dict, Any, List
4from pymongo import MongoClient, ReturnDocument
5from pymongo.errors import DuplicateKeyError, ConnectionFailure
6import bcrypt
7from dotenv import load_dotenv
8import logging
9
10# Configure logging
11logging.basicConfig(level=logging.INFO)
12logger = logging.getLogger(__name__)
13
14load_dotenv()
15
16class Database:
17    def __init__(self):
18        """Initialize database connection and ensure indexes"""
19        try:
20            mongo_uri = os.getenv("MONGO_URI")
21            if not mongo_uri:
22                raise ValueError("MONGO_URI environment variable is not set")
23                
24            self.client = MongoClient(
25                mongo_uri,
26                serverSelectionTimeoutMS=5000,  # 5 second timeout
27                connectTimeoutMS=30000,        # 30 second connection timeout
28                socketTimeoutMS=45000,          # 45 second socket timeout
29                connect=False                   # Lazy connection
30            )
31            
32            # Test the connection
33            self.client.admin.command('ping')
34            
35            self.db = self.client[os.getenv("DB_NAME", "rag_system")]
36            self.users = self.db["users"]
37            self._create_indexes()
38            logger.info("Successfully connected to MongoDB")
39            
40        except Exception as e:
41            logger.error(f"Failed to connect to MongoDB: {str(e)}")
42            raise
43    
44    def _create_indexes(self):
45        """Create necessary database indexes"""
46        try:
47            # Create unique index on username
48            self.users.create_index("username", unique=True)
49            logger.info("Created database indexes")
50        except Exception as e:
51            logger.error(f"Error creating indexes: {str(e)}")
52            raise
53    
54    def add_user(self, username: str, password: str, role: str) -> bool:
55        """Add a new user to the database"""
56        if not username or not password or not role:
57            logger.warning("Missing required fields for user creation")
58            return False
59            
60        try:
61            hashed = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt())
62            user_data = {
63                "username": username,
64                "password": hashed.decode('utf-8'),
65                "role": role.lower(),
66                "created_at": datetime.utcnow(),
67                "last_login": None
68            }
69            
70            result = self.users.insert_one(user_data)
71            if result.inserted_id:
72                logger.info(f"Created new user: {username}")
73                return True
74            return False
75            
76        except DuplicateKeyError:
77            logger.warning(f"Username already exists: {username}")
78            return False
79        except Exception as e:
80            logger.error(f"Error adding user {username}: {str(e)}")
81            return False
82    
83    def verify_user(self, username: str, password: str) -> Optional[Dict[str, Any]]:
84        """Verify user credentials"""
85        try:
86            user = self.users.find_one({"username": username})
87            if not user:
88                logger.warning(f"Login attempt for non-existent user: {username}")
89                return None
90                
91            if bcrypt.checkpw(password.encode('utf-8'), user["password"].encode('utf-8')):
92                # Update last login time
93                self.users.update_one(
94                    {"_id": user["_id"]},
95                    {"$set": {"last_login": datetime.utcnow()}}
96                )
97                logger.info(f"Successful login for user: {username}")
98                return {
99                    "username": user["username"],
100                    "role": user["role"],
101                    "last_login": user.get("last_login")
102                }
103            
104            logger.warning(f"Failed login attempt for user: {username}")
105            return None
106            
107        except Exception as e:
108            logger.error(f"Error verifying user {username}: {str(e)}")
109            return None
110    
111    def get_user(self, username: str) -> Optional[Dict[str, Any]]:
112        """Get user by username (without sensitive data)"""
113        try:
114            user = self.users.find_one(
115                {"username": username},
116                {"password": 0}  # Exclude password from results
117            )
118            return user
119        except Exception as e:
120            logger.error(f"Error fetching user {username}: {str(e)}")
121            return None
122
123# Initialize database connection
124db = Database()
125
126def initialize_users():
127    """
128    Initialize default users if they don't exist.
129    Returns tuple of (success_count, total_users, errors)
130    """
131    from datetime import datetime
132    
133    default_users = [
134        {"username": "Tony", "password": "password123", "role": "engineering"},
135        {"username": "Bruce", "password": "securepass", "role": "marketing"},
136        {"username": "Sam", "password": "financepass", "role": "finance"},
137        {"username": "Peter", "password": "pete123", "role": "engineering"},
138        {"username": "Sid", "password": "sidpass123", "role": "marketing"},
139        {"username": "Natasha", "password": "hrpass123", "role": "hr"}
140    ]
141    
142    success_count = 0
143    errors = []
144    
145    for user in default_users:
146        try:
147            if db.add_user(user["username"], user["password"], user["role"]):
148                success_count += 1
149                logger.info(f"Initialized user: {user['username']}")
150            else:
151                errors.append(f"Failed to add user: {user['username']}")
152        except Exception as e:
153            error_msg = f"Error initializing user {user['username']}: {str(e)}"
154            logger.error(error_msg)
155            errors.append(error_msg)
156    
157    logger.info(f"User initialization complete. Success: {success_count}/{len(default_users)}")
158    if errors:
159        logger.warning(f"Encountered {len(errors)} errors during user initialization")
160    
161    return success_count, len(default_users), errors
162