Refactor bot connecting with button on profile page

Munoon · Jun 6, 2020 · ac2eb2e · ac2eb2e
1 parent a148f00
commit ac2eb2e
Show file tree

Hide file tree

Showing 14 changed files with 166 additions and 68 deletions.
diff --git a/src/main/java/com/habr/telegrambotmfa/AuthorizedUser.java b/src/main/java/com/habr/telegrambotmfa/AuthorizedUser.java
@@ -13,4 +13,8 @@ public AuthorizedUser(User user) {
     public User getUser() {
         return user;
     }
+
+    public int getUserId() {
+        return user.getId();
+    }
 }
diff --git a/src/main/java/com/habr/telegrambotmfa/TelegramBot.java b/src/main/java/com/habr/telegrambotmfa/TelegramBot.java
@@ -1,7 +1,7 @@
 package com.habr.telegrambotmfa;
 
-import com.habr.telegrambotmfa.botCommands.ConnectAccountCommand;
-import com.habr.telegrambotmfa.botCommands.MfaCommand;
+import com.habr.telegrambotmfa.login.MfaCommand;
+import com.habr.telegrambotmfa.repositories.ConnectTelegramRepository;
 import org.springframework.context.annotation.Lazy;
 import org.springframework.core.env.Environment;
 import org.springframework.stereotype.Component;
@@ -21,14 +21,14 @@ public class TelegramBot extends TelegramLongPollingCommandBot {
     private MfaCommand mfaCommand;
     private String botUsername;
     private String botToken;
+    private ConnectTelegramRepository connectTelegramRepository;
 
-    public TelegramBot(Environment env, ConnectAccountCommand connectAccountCommand, @Lazy MfaCommand mfaCommand) throws TelegramApiException {
+    public TelegramBot(Environment env, @Lazy MfaCommand mfaCommand, ConnectTelegramRepository connectTelegramRepository) throws TelegramApiException {
         super(ApiContext.getInstance(DefaultBotOptions.class), false);
         this.botToken = env.getRequiredProperty("telegram.bot.token");
         this.mfaCommand = mfaCommand;
         this.botUsername = getMe().getUserName();
-
-        register(connectAccountCommand);
+        this.connectTelegramRepository = connectTelegramRepository;
     }
 
     @PostConstruct
@@ -41,6 +41,11 @@ public void addBot() throws TelegramApiRequestException {
     public void processNonCommandUpdate(Update update) {
         if (update.hasCallbackQuery()) {
             mfaCommand.onCallbackQuery(update.getCallbackQuery());
+        } else if (update.hasMessage()) {
+            var message = update.getMessage();
+            var telegramUserId = message.getFrom().getId();
+            var chatId = message.getChatId();
+            connectTelegramRepository.register(telegramUserId, chatId);
         }
     }
 

diff --git a/src/main/java/com/habr/telegrambotmfa/botCommands/ConnectAccountCommand.java b/src/main/java/com/habr/telegrambotmfa/botCommands/ConnectAccountCommand.java
diff --git a/src/main/java/com/habr/telegrambotmfa/config/WebSecurityConfig.java b/src/main/java/com/habr/telegrambotmfa/config/WebSecurityConfig.java
@@ -1,6 +1,6 @@
 package com.habr.telegrambotmfa.config;
 
-import com.habr.telegrambotmfa.botCommands.MfaCommand;
+import com.habr.telegrambotmfa.login.MfaCommand;
 import com.habr.telegrambotmfa.login.CustomAuthenticationProvider;
 import com.habr.telegrambotmfa.login.CustomFailureHandler;
 import com.habr.telegrambotmfa.login.CustomSuccessHandler;

diff --git a/src/main/java/com/habr/telegrambotmfa/controllers/UserController.java b/src/main/java/com/habr/telegrambotmfa/controllers/UserController.java
@@ -1,13 +1,18 @@
 package com.habr.telegrambotmfa.controllers;
 
+import com.habr.telegrambotmfa.AuthorizedUser;
+import com.habr.telegrambotmfa.models.ConnectTelegramTo;
 import com.habr.telegrambotmfa.models.User;
+import com.habr.telegrambotmfa.services.TelegramService;
 import com.habr.telegrambotmfa.services.UserService;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.security.core.annotation.AuthenticationPrincipal;
 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 org.telegram.telegrambots.meta.exceptions.TelegramApiException;
 
 import java.util.Set;
 
@@ -16,9 +21,11 @@
 public class UserController {
     private static final Logger log = LoggerFactory.getLogger(UserController.class);
     private UserService userService;
+    private TelegramService telegramService;
 
-    public UserController(UserService userService) {
+    public UserController(UserService userService, TelegramService telegramService) {
         this.userService = userService;
+        this.telegramService = telegramService;
     }
 
     @PostMapping("/register")
@@ -27,4 +34,14 @@ public User registerUser(@RequestBody User user) {
         user.setRoles(Set.of(User.Role.ROLE_USER));
         return userService.create(user);
     }
+
+    @PostMapping("/telegram")
+    public void connectTelegram(@RequestBody ConnectTelegramTo connectTelegram,
+                                @AuthenticationPrincipal AuthorizedUser authUser) throws TelegramApiException {
+        long chatId = telegramService.connectBot(authUser.getUserId(), connectTelegram.getUserId());
+
+        String userName = authUser.getUser().getName();
+        String welcomeMessage = String.format("Добро пожаловать, %s!\nВы успешно подключили аккаунт!", userName);
+        telegramService.sendMessage(welcomeMessage, chatId);
+    }
 }
diff --git a/src/main/java/com/habr/telegrambotmfa/login/CustomAuthenticationProvider.java b/src/main/java/com/habr/telegrambotmfa/login/CustomAuthenticationProvider.java
@@ -1,7 +1,6 @@
 package com.habr.telegrambotmfa.login;
 
 import com.habr.telegrambotmfa.AuthorizedUser;
-import com.habr.telegrambotmfa.botCommands.MfaCommand;
 import com.habr.telegrambotmfa.models.User;
 import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
 import org.springframework.security.authentication.dao.DaoAuthenticationProvider;

diff --git a/...elegrambotmfa/botCommands/MfaCommand.java → ...habr/telegrambotmfa/login/MfaCommand.java b/...elegrambotmfa/botCommands/MfaCommand.java → ...habr/telegrambotmfa/login/MfaCommand.java
@@ -1,9 +1,7 @@
-package com.habr.telegrambotmfa.botCommands;
+package com.habr.telegrambotmfa.login;
 
 import com.habr.telegrambotmfa.AuthorizedUser;
 import com.habr.telegrambotmfa.TelegramBot;
-import com.habr.telegrambotmfa.login.AuthenticationInfo;
-import com.habr.telegrambotmfa.login.CustomSuccessHandler;
 import com.habr.telegrambotmfa.models.User;
 import com.habr.telegrambotmfa.services.WebSocketService;
 import org.slf4j.Logger;
@@ -17,7 +15,6 @@
 import org.telegram.telegrambots.meta.api.methods.updatingmessages.EditMessageText;
 import org.telegram.telegrambots.meta.api.objects.CallbackQuery;
 import org.telegram.telegrambots.meta.api.objects.Message;
-import org.telegram.telegrambots.meta.api.objects.Update;
 import org.telegram.telegrambots.meta.api.objects.replykeyboard.InlineKeyboardMarkup;
 import org.telegram.telegrambots.meta.api.objects.replykeyboard.buttons.InlineKeyboardButton;
 import org.telegram.telegrambots.meta.exceptions.TelegramApiException;

diff --git a/src/main/java/com/habr/telegrambotmfa/models/ConnectTelegramTo.java b/src/main/java/com/habr/telegrambotmfa/models/ConnectTelegramTo.java
@@ -0,0 +1,25 @@
+package com.habr.telegrambotmfa.models;
+
+public class ConnectTelegramTo {
+    private String checkString;
+    private int userId;
+
+    public ConnectTelegramTo() {
+    }
+
+    public String getCheckString() {
+        return checkString;
+    }
+
+    public void setCheckString(String checkString) {
+        this.checkString = checkString;
+    }
+
+    public int getUserId() {
+        return userId;
+    }
+
+    public void setUserId(int userId) {
+        this.userId = userId;
+    }
+}
diff --git a/src/main/java/com/habr/telegrambotmfa/repositories/ConnectTelegramRepository.java b/src/main/java/com/habr/telegrambotmfa/repositories/ConnectTelegramRepository.java
@@ -0,0 +1,23 @@
+package com.habr.telegrambotmfa.repositories;
+
+import org.springframework.stereotype.Repository;
+
+import java.util.HashMap;
+import java.util.Map;
+
+@Repository
+public class ConnectTelegramRepository {
+    private Map<Integer, Long> repository = new HashMap<>(); // user id, chat id
+
+    public void register(int userId, long chatId) {
+        repository.put(userId, chatId);
+    }
+
+    public Long getChatIdByUserId(int userId) {
+        return repository.get(userId);
+    }
+
+    public void removeByUserId(int userId) {
+        repository.remove(userId);
+    }
+}
diff --git a/src/main/java/com/habr/telegrambotmfa/repositories/UserRepository.java b/src/main/java/com/habr/telegrambotmfa/repositories/UserRepository.java
@@ -14,6 +14,6 @@ public interface UserRepository extends JpaRepository<User, Integer> {
 
     @Modifying
     @Transactional
-    @Query("UPDATE User u SET u.telegramChatId = ?2 WHERE u.username = ?1")
-    void connectBot(String username, Long chatId);
+    @Query("UPDATE User u SET u.telegramChatId = ?2 WHERE u.id = ?1")
+    void connectBot(int userId, Long chatId);
 }
diff --git a/src/main/java/com/habr/telegrambotmfa/services/TelegramService.java b/src/main/java/com/habr/telegrambotmfa/services/TelegramService.java
@@ -0,0 +1,37 @@
+package com.habr.telegrambotmfa.services;
+
+import com.habr.telegrambotmfa.TelegramBot;
+import com.habr.telegrambotmfa.repositories.ConnectTelegramRepository;
+import com.habr.telegrambotmfa.repositories.UserRepository;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.stereotype.Service;
+import org.telegram.telegrambots.meta.api.methods.send.SendMessage;
+import org.telegram.telegrambots.meta.exceptions.TelegramApiException;
+
+@Service
+public class TelegramService {
+    private UserRepository userRepository;
+    private ConnectTelegramRepository connectTelegramRepository;
+    private TelegramBot telegramBot;
+
+    @Autowired
+    public TelegramService(UserRepository userRepository, ConnectTelegramRepository connectTelegramRepository, TelegramBot telegramBot) {
+        this.userRepository = userRepository;
+        this.connectTelegramRepository = connectTelegramRepository;
+        this.telegramBot = telegramBot;
+    }
+
+    public long connectBot(int userId, int telegramUserId) {
+        long chatId = connectTelegramRepository.getChatIdByUserId(telegramUserId);
+        userRepository.connectBot(userId, chatId);
+        connectTelegramRepository.removeByUserId(telegramUserId);
+        return chatId;
+    }
+
+    public void sendMessage(String text, long chatId) throws TelegramApiException {
+        var message = new SendMessage()
+                .setText(text)
+                .setChatId(chatId);
+        telegramBot.execute(message);
+    }
+}
diff --git a/src/main/java/com/habr/telegrambotmfa/services/UserService.java b/src/main/java/com/habr/telegrambotmfa/services/UserService.java
@@ -24,10 +24,6 @@ public UserService(UserRepository userRepository, @Lazy PasswordEncoder password
         this.passwordEncoder = passwordEncoder;
     }
 
-    public void connectBot(String username, Long chatId) {
-        userRepository.connectBot(username, chatId);
-    }
-
     public User create(User user) {
         user.setId(null);
         user.setPassword(passwordEncoder.encode(user.getPassword()));

diff --git a/src/main/resources/application.properties b/src/main/resources/application.properties
@@ -1,3 +1,5 @@
+server.port=80
+
 spring.datasource.platform=h2
 spring.datasource.url=jdbc:h2:mem:telegram-mfa
 spring.datasource.driverClassName=org.h2.Driver

diff --git a/src/main/resources/templates/index.html b/src/main/resources/templates/index.html
@@ -4,8 +4,13 @@
     <meta charset="UTF-8">
     <title>Telegram Bot MFA</title>
 
+    <meta name="_csrf" th:content="${_csrf.token}"/>
+    <meta name="_csrf_header" th:content="${_csrf.headerName}"/>
+
     <link rel="stylesheet" th:href="@{/webjars/bootstrap/4.4.1-1/css/bootstrap.min.css}"/>
     <link rel="stylesheet" th:href="@{/static/style.css}"/>
+
+    <script th:src="@{/webjars/jquery/3.4.0/jquery.min.js}" defer></script>
 </head>
 <body>
 <div class="container">
@@ -22,18 +27,44 @@ <h5 class="card-title text-center">Главная страница</h5>
                         <span th:if="${useTelegram && telegram == null}">Ваш телеграм аккаунт <b>подключён</b>.</span>
                     </div>
 
-                    <div class="mt-3" th:if="${!useTelegram}" th:with="botName = ${@telegramBot.getBotUsername()}">
-                        <hr>
-                        <span>
-                            Для подключения Telegram аккаунта, откройте бота
-                            <a th:href="${'https://t.me/' + botName}">@[[${botName}]]</a>,
-                            запустите его командой <b>/start</b> и введите команду
-                            <b>/connect [[${#authentication.principal.user.username}]]</b>.
-                            Затем обновите эту страницу, что бы убедиться, что вы подключили бота.
-                            Вы должны увидеть свой ник в Telegram или сообщение о том, что ваш Telegram аккаунт подключен.
-                        </span>
-                        <hr>
-                    </div>
+                    <th:block th:if="${!useTelegram}">
+                        <script async src="https://telegram.org/js/telegram-widget.js?9"
+                                data-telegram-login="MunoonBot_bot"
+                                data-size="large" data-onauth="onTelegramAuth(user)"
+                                data-request-access="write"></script>
+
+                        <script>
+                            function onTelegramAuth(user) {
+                                $.post({
+                                    url: '/ajax/user/telegram',
+                                    method: 'POST',
+                                    contentType: 'application/json',
+                                    headers: getHeaders(),
+                                    data: JSON.stringify({
+                                        checkString: getCheckString(user),
+                                        userId: user.id
+                                    })
+                                });
+                            }
+
+                            function getCheckString(data) {
+                                let result = '';
+                                for (let item in data) {
+                                    if (data.hasOwnProperty(item)) {
+                                        result += `${item}=${data[item]}\n`;
+                                    }
+                                }
+                                return result;
+                            }
+
+                            function getHeaders() {
+                                let headers = {};
+                                let csrfHeader = document.querySelector('meta[name="_csrf_header"]').content;
+                                headers[csrfHeader] = document.querySelector('meta[name="_csrf"]').content;
+                                return headers;
+                            }
+                        </script>
+                    </th:block>
 
                     <form th:action="@{/logout}" method="post" class="mt-3">
                         <button class="btn btn-secondary" type="submit">Выйти</button>