import json
import base64
import os
from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from getpass import getpass  # Untuk input password yang lebih aman

class PasswordManager:
    def __init__(self):
        self.data_file = "passwords.enc"
        self.salt_file = "salt.key"
        self.fernet = None

    def generate_key(self, master_password):
        """Membuat kunci enkripsi dari master password"""
        # Coba muat salt yang ada atau buat yang baru
        try:
            with open(self.salt_file, "rb") as f:
                salt = f.read()
        except FileNotFoundError:
            salt = os.urandom(16)
            with open(self.salt_file, "wb") as f:
                f.write(salt)
    
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=600000,  # Meningkatkan iterasi untuk keamanan lebih
        )
        key = base64.urlsafe_b64encode(kdf.derive(master_password.encode()))
        self.fernet = Fernet(key)

    def load_data(self):
        """Memuat data yang terenkripsi"""
        if not self.fernet:
            raise ValueError("Fernet key not initialized")
        
        try:
            with open(self.data_file, "rb") as f:
                encrypted_data = f.read()
                decrypted_data = self.fernet.decrypt(encrypted_data)
                return json.loads(decrypted_data)
        except FileNotFoundError:
            return {}
        except InvalidToken:
            print("Master password salah atau file korup!")
            return None
        except Exception as e:
            print(f"Error loading data: {str(e)}")
            return None

    def save_data(self, data):
        """Menyimpan data terenkripsi"""
        if not self.fernet:
            raise ValueError("Fernet key not initialized")
        
        try:
            json_data = json.dumps(data)
            encrypted_data = self.fernet.encrypt(json_data.encode())
        
            # Buat backup sebelum menimpa file
            if os.path.exists(self.data_file):
                os.replace(self.data_file, self.data_file + ".bak")
        
            with open(self.data_file, "wb") as f:
                f.write(encrypted_data)
        
            # Hapus backup jika berhasil
            if os.path.exists(self.data_file + ".bak"):
                os.remove(self.data_file + ".bak")
            
            return True
        except Exception as e:
            print(f"Error saving data: {str(e)}")
            # Restore backup jika ada
            if os.path.exists(self.data_file + ".bak"):
                os.replace(self.data_file + ".bak", self.data_file)
            return False

    def add_account(self, service, username, password):
        """Menambahkan akun baru"""
        data = self.load_data()
        if data is None:  # Jika error decrypting
            return False
        
        if not data:  # Jika data kosong
            data = {}
        
        data[service] = {
            'username': username,
            'password': password
        }
        return self.save_data(data)

    def get_password(self, service):
        """Mendapatkan password untuk layanan tertentu"""
        data = self.load_data()
        return data.get(service) if data else None

    def list_services(self):
        """Menampilkan daftar semua layanan yang tersimpan"""
        data = self.load_data()
        return list(data.keys()) if data else []

    def delete_account(self, service):
        """Menghapus akun"""
        data = self.load_data()
        if data is None:  # Jika error decrypting
            return False
        
        if data and service in data:
            del data[service]
            return self.save_data(data)
        return False


def main():
    manager = PasswordManager()

    # Setup master password
    while True:
        master_pass = getpass("Masukkan master password: ").strip()
        if not master_pass:
            print("Password tidak boleh kosong!")
            continue
        
        try:
            manager.generate_key(master_pass)
            # Test decrypt untuk memverifikasi password
            test_data = manager.load_data()
            if test_data is None and os.path.exists(manager.data_file):
                # Jika file ada tapi decrypt gagal, password salah
                print("Master password salah!")
                continue
            break
        except Exception as e:
            print(f"Error: {str(e)}")
            return
  
    while True:
        print("\n=== Password Manager ===")
        print("1. Tambah Password Baru")
        print("2. Lihat Password")
        print("3. List Semua Akun")
        print("4. Hapus Akun")
        print("5. Keluar")
        
        choice = input("Pilihan (1-5): ").strip()
    
        if choice == "1":
            service = input("Nama Layanan (contoh: Facebook/Gmail): ").strip()
            if not service:
                print("Nama layanan tidak boleh kosong!")
                continue
            
            username = input("Username/Email: ").strip()
            password = getpass("Password: ").strip()
        
            if manager.add_account(service, username, password):
                print("Berhasil menyimpan password!")
            else:
                print("Gagal menyimpan password")
    
        elif choice == "2":
            service = input("Nama Layanan: ").strip()
            account = manager.get_password(service)
            if account:
                print(f"\nService: {service}")
                print(f"Username: {account['username']}")
                print(f"Password: {account['password']}")
            else:
                print("Akun tidak ditemukan")
    
        elif choice == "3":
            services = manager.list_services()
            if services:
                print("\nDaftar Layanan:")
                for i, service in enumerate(services, 1):
                    print(f"{i}. {service}")
            else:
                print("Tidak ada akun yang tersimpan")
    
        elif choice == "4":
            service = input("Nama Layanan yang akan dihapus: ").strip()
            if manager.delete_account(service):
                print("Akun berhasil dihapus")
            else:
                print("Gagal menghapus akun atau akun tidak ditemukan")
    
        elif choice == "5":
            print("Keluar dari aplikasi...")
            break
    
        else:
            print("Pilihan tidak valid!")

if __name__ == "__main__":
    main()