ssh-manager/backend/src/main/java/com/sshmanager/controller/AuthController.java

package com.sshmanager.controller;

import com.sshmanager.dto.LoginRequest;
import com.sshmanager.dto.LoginResponse;
import com.sshmanager.dto.ChangePasswordRequest;
import com.sshmanager.dto.RegisterRequest;
import com.sshmanager.dto.UserDto;
import com.sshmanager.exception.InvalidOperationException;
import com.sshmanager.service.UserService;
import com.sshmanager.entity.User;
import com.sshmanager.repository.UserRepository;
import com.sshmanager.security.JwtTokenProvider;
import org.springframework.http.ResponseEntity;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.authentication.BadCredentialsException;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.crypto.password.PasswordEncoder;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.time.Instant;
import java.util.HashMap;
import java.util.Map;

@RestController
@RequestMapping("/api/auth")
public class AuthController {

    private final AuthenticationManager authenticationManager;
    private final JwtTokenProvider tokenProvider;
    private final UserRepository userRepository;
    private final PasswordEncoder passwordEncoder;
    private final UserService userService;

    public AuthController(AuthenticationManager authenticationManager,
                          JwtTokenProvider tokenProvider,
                          UserRepository userRepository,
                          PasswordEncoder passwordEncoder,
                          UserService userService) {
        this.authenticationManager = authenticationManager;
        this.tokenProvider = tokenProvider;
        this.userRepository = userRepository;
        this.passwordEncoder = passwordEncoder;
        this.userService = userService;
    }

    @PostMapping("/login")
    public ResponseEntity<?> login(@RequestBody LoginRequest request) {
        try {
            Authentication authentication = authenticationManager.authenticate(
                new UsernamePasswordAuthenticationToken(request.getUsername(), request.getPassword()));

            SecurityContextHolder.getContext().setAuthentication(authentication);
            String token = tokenProvider.generateToken(authentication);

            User user = userRepository.findByUsername(request.getUsername()).orElseThrow(() -> new IllegalStateException("User not found"));
            LoginResponse response = new LoginResponse(token, user.getUsername(),
                user.getDisplayName() != null ? user.getDisplayName() : user.getUsername(),
                user.getRole() != null ? user.getRole() : "ROLE_USER",
                isPasswordChangeRequired(user));

            return ResponseEntity.ok(response);
        } catch (BadCredentialsException e) {
            Map<String, String> error = new HashMap<>();
            error.put("message", "Invalid username or password");
            return ResponseEntity.status(401).body(error);
        }
    }

    @PostMapping("/register")
    public ResponseEntity<?> register(@RequestBody RegisterRequest request) {
        try {
            UserDto created = userService.createRegularUser(
                    request.getUsername(), request.getPassword(), request.getDisplayName());

            // Auto-login: construct authentication directly (user already created & verified)
            UsernamePasswordAuthenticationToken auth =
                    new UsernamePasswordAuthenticationToken(created.getUsername(), null);
            SecurityContextHolder.getContext().setAuthentication(auth);
            String token = tokenProvider.generateToken(auth);

            LoginResponse response = new LoginResponse(token, created.getUsername(),
                    created.getDisplayName() != null ? created.getDisplayName() : created.getUsername(),
                    created.getRole() != null ? created.getRole() : "ROLE_USER",
                    false);

            return ResponseEntity.ok(response);
        } catch (InvalidOperationException e) {
            Map<String, String> error = new HashMap<>();
            error.put("message", e.getMessage());
            return ResponseEntity.badRequest().body(error);
        }
    }

    @GetMapping("/health")
    public ResponseEntity<?> health() {
        Map<String, Object> data = new HashMap<>();
        data.put("app", "ssh-manager");
        data.put("status", "ok");
        data.put("timestamp", Instant.now().toEpochMilli());
        return ResponseEntity.ok(data);
    }

    @GetMapping("/me")
    public ResponseEntity<?> me(Authentication authentication) {
        if (authentication == null || !authentication.isAuthenticated()) {
            Map<String, String> error = new HashMap<>();
            error.put("error", "Unauthorized");
            return ResponseEntity.status(401).body(error);
        }
        User user = userRepository.findByUsername(authentication.getName()).orElseThrow(() -> new IllegalStateException("User not found"));
        Map<String, Object> data = new HashMap<>();
        data.put("username", user.getUsername());
        data.put("displayName", user.getDisplayName());
        data.put("role", user.getRole() != null ? user.getRole() : "ROLE_USER");
        data.put("passwordChangeRequired", isPasswordChangeRequired(user));
        return ResponseEntity.ok(data);
    }

    @PostMapping("/change-password")
    public ResponseEntity<?> changePassword(@RequestBody ChangePasswordRequest request, Authentication authentication) {
        if (authentication == null || !authentication.isAuthenticated()) {
            Map<String, String> error = new HashMap<>();
            error.put("error", "Unauthorized");
            return ResponseEntity.status(401).body(error);
        }

        String currentPassword = request.getCurrentPassword() == null ? "" : request.getCurrentPassword().trim();
        String newPassword = request.getNewPassword() == null ? "" : request.getNewPassword().trim();

        if (currentPassword.isEmpty() || newPassword.isEmpty()) {
            Map<String, String> error = new HashMap<>();
            error.put("message", "Current password and new password are required");
            return ResponseEntity.badRequest().body(error);
        }

        if (newPassword.length() < 8) {
            Map<String, String> error = new HashMap<>();
            error.put("message", "New password must be at least 8 characters");
            return ResponseEntity.badRequest().body(error);
        }

        User user = userRepository.findByUsername(authentication.getName()).orElseThrow(() -> new IllegalStateException("User not found"));
        if (!passwordEncoder.matches(currentPassword, user.getPasswordHash())) {
            Map<String, String> error = new HashMap<>();
            error.put("message", "Current password is incorrect");
            return ResponseEntity.badRequest().body(error);
        }

        if (passwordEncoder.matches(newPassword, user.getPasswordHash())) {
            Map<String, String> error = new HashMap<>();
            error.put("message", "New password must be different from current password");
            return ResponseEntity.badRequest().body(error);
        }

        user.setPasswordHash(passwordEncoder.encode(newPassword));
        user.setPasswordChangedAt(Instant.now());
        userRepository.save(user);

        Map<String, Object> data = new HashMap<>();
        data.put("message", "Password updated");
        data.put("passwordChangeRequired", false);
        return ResponseEntity.ok(data);
    }

    private boolean isPasswordChangeRequired(User user) {
        return user.getPasswordChangedAt() == null
            || !user.getPasswordChangedAt().isAfter(Instant.EPOCH);
    }
}