app.py

# Copyright 2025 DBMS Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from flask import Flask, render_template, request, jsonify, redirect, url_for, flash, session
import boto3
import os
from werkzeug.utils import secure_filename
from dotenv import load_dotenv
from datetime import datetime
import mimetypes
from urllib.parse import quote
import hashlib

app = Flask(__name__)
app.secret_key = os.getenv('JWT_SECRET', 'fallback_secret_key')

# Load environment variables
load_dotenv()

# Login credentials
LOGIN_USERNAME = "theops"
LOGIN_PASSWORD = "theopss"

def login_required(f):
    """Decorator to require login for protected routes"""
    def decorated_function(*args, **kwargs):
        if not session.get('logged_in'):
            return redirect(url_for('login'))
        return f(*args, **kwargs)
    decorated_function.__name__ = f.__name__
    return decorated_function

# AWS S3 Configuration
AWS_ACCESS_KEY_ID = os.getenv('AWS_ACCESS_KEY_ID')
AWS_SECRET_ACCESS_KEY = os.getenv('AWS_SECRET_ACCESS_KEY')
AWS_DEFAULT_REGION = os.getenv('AWS_DEFAULT_REGION')
S3_BUCKET_NAME = os.getenv('S3_BUCKET_NAME')

# Initialize S3 client
s3_client = boto3.client(
    's3',
    aws_access_key_id=AWS_ACCESS_KEY_ID,
    aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
    region_name=AWS_DEFAULT_REGION
)

# Allowed file extensions
ALLOWED_EXTENSIONS = {
    'txt', 'pdf', 'png', 'jpg', 'jpeg', 'gif', 'doc', 'docx', 
    'xls', 'xlsx', 'ppt', 'pptx', 'csv', 'zip', 'rar', '7z',
    'mp4', 'avi', 'mov', 'mp3', 'wav', 'json', 'xml'
}

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def get_s3_folders_and_files():
    """Get all folders with their sizes from S3 bucket"""
    try:
        folders_data = {}
        files = []
        
        # Use paginator to handle large buckets
        paginator = s3_client.get_paginator('list_objects_v2')
        pages = paginator.paginate(Bucket=S3_BUCKET_NAME)
        
        for page in pages:
            if 'Contents' in page:
                for obj in page['Contents']:
                    key = obj['Key']
                    
                    # Skip if it's a folder marker (ends with /)
                    if key.endswith('/'):
                        continue
                    
                    # Extract folder path
                    if '/' in key:
                        folder_path = '/'.join(key.split('/')[:-1])
                        
                        # Initialize folder data if not exists
                        if folder_path not in folders_data:
                            folders_data[folder_path] = {
                                'name': folder_path,
                                'size': 0,
                                'file_count': 0,
                                'last_modified': obj['LastModified']
                            }
                        
                        # Add file size to folder
                        folders_data[folder_path]['size'] += obj['Size']
                        folders_data[folder_path]['file_count'] += 1
                        
                        # Update last modified if this file is newer
                        if obj['LastModified'] > folders_data[folder_path]['last_modified']:
                            folders_data[folder_path]['last_modified'] = obj['LastModified']
                    else:
                        # Root level files
                        if 'Root' not in folders_data:
                            folders_data['Root'] = {
                                'name': 'Root',
                                'size': 0,
                                'file_count': 0,
                                'last_modified': obj['LastModified']
                            }
                        
                        folders_data['Root']['size'] += obj['Size']
                        folders_data['Root']['file_count'] += 1
                        
                        if obj['LastModified'] > folders_data['Root']['last_modified']:
                            folders_data['Root']['last_modified'] = obj['LastModified']
                    
                    # Add file info for individual file operations
                    files.append({
                        'key': key,
                        'filename': key.split('/')[-1],
                        'folder': '/'.join(key.split('/')[:-1]) if '/' in key else 'Root',
                        'size': obj['Size'],
                        'last_modified': obj['LastModified'],
                        'url': f"https://{S3_BUCKET_NAME}.s3.{AWS_DEFAULT_REGION}.amazonaws.com/{quote(key)}"
                    })
        
        # Convert folders_data to list and sort
        folders = list(folders_data.values())
        folders.sort(key=lambda x: x['name'])
        
        return folders, files
    except Exception as e:
        print(f"Error getting S3 data: {str(e)}")
        return [], []

@app.route('/login', methods=['GET', 'POST'])
def login():
    """Login page"""
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        
        if username == LOGIN_USERNAME and password == LOGIN_PASSWORD:
            session['logged_in'] = True
            session['username'] = username
            flash('Login successful!', 'success')
            return redirect(url_for('upload_page'))
        else:
            flash('Invalid username or password', 'error')
    
    return render_template('login.html')

@app.route('/logout')
def logout():
    """Logout and clear session"""
    session.clear()
    flash('You have been logged out', 'info')
    return redirect(url_for('login'))

@app.route('/')
@login_required
def index():
    """Main dashboard page with navigation to upload and folders sections"""
    folders, files = get_s3_folders_and_files()
    return render_template('index.html', folders=folders, files=files)

@app.route('/upload')
@login_required
def upload_page():
    """Upload page for file and folder uploads"""
    folders, files = get_s3_folders_and_files()
    return render_template('upload.html', folders=folders, files=files)

@app.route('/folders')
@login_required
def folders_page():
    """Folders page showing available folders and files"""
    folders, files = get_s3_folders_and_files()
    search_query = request.args.get('search', '')
    
    # Filter folders based on search query
    if search_query:
        folders = [f for f in folders if search_query.lower() in f['name'].lower()]
    
    return render_template('folders.html', folders=folders, files=files, search_query=search_query)

@app.route('/upload', methods=['POST'])
@login_required
def upload_files():
    """Handle file uploads with folder structure"""
    try:
        if 'files' not in request.files:
            flash('No files selected', 'error')
            return redirect(url_for('upload_page'))
        
        files = request.files.getlist('files')
        target_folder = request.form.get('target_folder', '').strip()
        
        if not files or all(file.filename == '' for file in files):
            flash('No files selected', 'error')
            return redirect(url_for('upload_page'))
        
        uploaded_count = 0
        errors = []
        
        for file in files:
            if file and file.filename != '':
                if allowed_file(file.filename):
                    try:
                        # Secure the filename
                        filename = secure_filename(file.filename)
                        
                        # Create S3 key with folder structure
                        if target_folder:
                            s3_key = f"{target_folder}/{filename}"
                        else:
                            s3_key = filename
                        
                        # Upload to S3
                        s3_client.upload_fileobj(
                            file,
                            S3_BUCKET_NAME,
                            s3_key,
                            ExtraArgs={
                                'ContentType': mimetypes.guess_type(filename)[0] or 'application/octet-stream'
                            }
                        )
                        uploaded_count += 1
                        
                    except Exception as e:
                        errors.append(f"Error uploading {file.filename}: {str(e)}")
                else:
                    errors.append(f"File type not allowed: {file.filename}")
        
        if uploaded_count > 0:
            flash(f'Successfully uploaded {uploaded_count} file(s)', 'success')
        
        if errors:
            for error in errors:
                flash(error, 'error')
        
        return redirect(url_for('upload_page'))
        
    except Exception as e:
        flash(f'Upload error: {str(e)}', 'error')
        return redirect(url_for('upload_page'))

def check_folder_exists(folder_name):
    """Check if a folder already exists in S3"""
    try:
        # Handle root folder case
        if folder_name == 'Root':
            # Check for files without folder prefix
            response = s3_client.list_objects_v2(
                Bucket=S3_BUCKET_NAME,
                Delimiter='/',
                MaxKeys=1
            )
            return 'Contents' in response and len(response['Contents']) > 0
        else:
            # List objects with the folder prefix
            response = s3_client.list_objects_v2(
                Bucket=S3_BUCKET_NAME,
                Prefix=f"{folder_name}/",
                MaxKeys=1
            )
            return 'Contents' in response and len(response['Contents']) > 0
    except Exception as e:
        print(f"Error checking folder existence: {str(e)}")
        return False

@app.route('/upload_folder', methods=['POST'])
@login_required
def upload_folder():
    """Handle folder uploads with all files inside"""
    try:
        if 'folder_files' not in request.files:
            flash('No folder selected', 'error')
            return redirect(url_for('upload_page'))
        
        files = request.files.getlist('folder_files')
        
        if not files or all(file.filename == '' for file in files):
            flash('No files in folder selected', 'error')
            return redirect(url_for('upload_page'))
        
        uploaded_count = 0
        skipped_count = 0
        errors = []
        folder_names = set()
        skipped_folders = []
        
        # First, collect all folder names from the files
        for file in files:
            if file and file.filename != '':
                relative_path = file.filename
                if '/' in relative_path:
                    folder_name = relative_path.split('/')[0]
                    folder_names.add(folder_name)
                else:
                    folder_names.add('Root')
        
        # Check which folders already exist
        existing_folders = set()
        for folder_name in folder_names:
            if check_folder_exists(folder_name):
                existing_folders.add(folder_name)
                skipped_folders.append(folder_name)
        
        # Process files, skipping those from existing folders
        for file in files:
            if file and file.filename != '':
                relative_path = file.filename
                
                # Determine which folder this file belongs to
                if '/' in relative_path:
                    folder_name = relative_path.split('/')[0]
                else:
                    folder_name = 'Root'
                
                # Skip if folder already exists
                if folder_name in existing_folders:
                    skipped_count += 1
                    continue
                
                # Check if file extension is allowed
                file_extension = relative_path.split('.')[-1].lower() if '.' in relative_path else ''
                if file_extension in ALLOWED_EXTENSIONS:
                    try:
                        # Create S3 key preserving the folder structure
                        s3_key = relative_path.replace('\\', '/')  # Handle Windows paths
                        
                        # Upload to S3
                        s3_client.upload_fileobj(
                            file,
                            S3_BUCKET_NAME,
                            s3_key,
                            ExtraArgs={
                                'ContentType': mimetypes.guess_type(relative_path)[0] or 'application/octet-stream'
                            }
                        )
                        uploaded_count += 1
                        
                    except Exception as e:
                        errors.append(f"Error uploading {relative_path}: {str(e)}")
                else:
                    errors.append(f"File type not allowed: {relative_path}")
        
        # Prepare success messages
        new_folders = folder_names - existing_folders
        
        if uploaded_count > 0:
            if new_folders:
                flash(f'✅ Successfully uploaded {len(new_folders)} new folder(s) with {uploaded_count} file(s)', 'success')
            else:
                flash(f'✅ Successfully uploaded {uploaded_count} file(s) to existing folders', 'success')
        
        if skipped_folders:
            if len(skipped_folders) == 1:
                flash(f'⏭️ Skipped existing folder: {skipped_folders[0]}', 'info')
            else:
                flash(f'⏭️ Skipped {len(skipped_folders)} existing folders: {", ".join(skipped_folders)}', 'info')
        
        if errors:
            for error in errors:
                flash(f'❌ {error}', 'error')
        
        if uploaded_count == 0 and skipped_count > 0:
            if len(skipped_folders) == 1:
                flash(f'ℹ️ Folder "{skipped_folders[0]}" already exists. No new files uploaded.', 'info')
            else:
                flash(f'ℹ️ All {len(skipped_folders)} folders already exist. No new files uploaded.', 'info')
        
        if uploaded_count == 0 and skipped_count == 0:
            flash('⚠️ No files were uploaded. Please check file types and try again.', 'error')
        
        return redirect(url_for('upload_page'))
        
    except Exception as e:
        flash(f'Folder upload error: {str(e)}', 'error')
        return redirect(url_for('upload_page'))

@app.route('/delete/<path:file_key>')
@login_required
def delete_file(file_key):
    """Delete a file from S3"""
    try:
        s3_client.delete_object(Bucket=S3_BUCKET_NAME, Key=file_key)
        flash(f'File {file_key} deleted successfully', 'success')
    except Exception as e:
        flash(f'Error deleting file: {str(e)}', 'error')
    
        return redirect(url_for('folders_page'))


@app.route('/api/folders')
@login_required
def api_folders():
    """API endpoint to get folders list"""
    folders, _ = get_s3_folders_and_files()
    return jsonify(folders)

@app.route('/api/files')
@login_required
def api_files():
    """API endpoint to get files list"""
    _, files = get_s3_folders_and_files()
    return jsonify(files)

@app.route('/api/upload_progress', methods=['POST'])
@login_required
def api_upload_progress():
    """API endpoint for real-time upload progress tracking"""
    try:
        data = request.get_json()
        folder_name = data.get('folder_name')
        progress = data.get('progress', 0)
        status = data.get('status', 'uploading')
        
        # Store progress in session or cache (for demo, we'll just return success)
        # In production, you might want to use Redis or database for this
        
        return jsonify({
            'success': True,
            'folder_name': folder_name,
            'progress': progress,
            'status': status
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)