Skip to content
Permalink
Browse files Browse the repository at this point in the history
[dy] Fix terminal user authentication (#2586)
* [dy] Update terminal user authentication

* [dy] Update auth
  • Loading branch information
dy46 committed May 4, 2023
1 parent e16a649 commit f63cd00
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 16 deletions.
34 changes: 22 additions & 12 deletions mage_ai/frontend/components/Terminal/index.tsx
@@ -1,6 +1,7 @@
import Ansi from 'ansi-to-react';
import { useCallback, useEffect, useMemo, useRef, useState } from 'react';

import AuthToken from '@api/utils/AuthToken';
import ClickOutside from '@oracle/components/ClickOutside';
import Text from '@oracle/elements/Text';
import {
Expand All @@ -26,6 +27,7 @@ import {
KEY_CODE_META,
KEY_CODE_V,
} from '@utils/hooks/keyboardShortcuts/constants';
import { OAUTH2_APPLICATION_CLIENT_ID } from '@api/constants';
import { onlyKeysPresent } from '@utils/hooks/keyboardShortcuts/utils';
import { pauseEvent } from '@utils/events';
import { useKeyboardContext } from '@context/Keyboard';
Expand Down Expand Up @@ -58,6 +60,14 @@ function Terminal({

const [stdout, setStdout] = useState<string>();

const token = useMemo(() => new AuthToken(), []);
const oauthWebsocketData = useMemo(() => ({
api_key: OAUTH2_APPLICATION_CLIENT_ID,
token: token.decodedToken.token,
}), [
token,
]);

useEffect(() => {
if (lastMessage) {
const msg = JSON.parse(lastMessage.data);
Expand Down Expand Up @@ -122,12 +132,10 @@ function Terminal({
}, [command]);

const sendCommand = useCallback((cmd) => {
sendMessage(JSON.stringify([
'stdin', cmd,
]));
sendMessage(JSON.stringify([
'stdin', '\r',
]));
sendMessage(JSON.stringify({
...oauthWebsocketData,
command: ['stdin', cmd + '\r'],
}));
if (cmd?.length >= 2) {
setCommandIndex(commandHistory.length + 1);
setCommandHistory(prev => prev.concat(cmd));
Expand Down Expand Up @@ -174,12 +182,14 @@ function Terminal({
pauseEvent(event);
if (onlyKeysPresent([KEY_CODE_CONTROL, KEY_CODE_C], keyMapping)) {
if (command?.length >= 0) {
sendMessage(JSON.stringify([
'stdin', command,
]));
sendMessage(JSON.stringify([
'stdin', '\x03',
]));
sendMessage(JSON.stringify({
...oauthWebsocketData,
command: ['stdin', command],
}));
sendMessage(JSON.stringify({
...oauthWebsocketData,
command: ['stdin', '\x03'],
}));
setCursorIndex(0);
}
setCommand('');
Expand Down
41 changes: 37 additions & 4 deletions mage_ai/server/terminal_server.py
@@ -1,6 +1,14 @@
from mage_ai.api.utils import authenticate_client_and_token
from mage_ai.api.utils import (
authenticate_client_and_token,
has_at_least_editor_role,
)
from mage_ai.orchestration.db.models.oauth import Oauth2Application
from mage_ai.settings import REQUIRE_USER_AUTHENTICATION
from mage_ai.settings import (
is_disable_pipeline_edit_access,
REQUIRE_USER_AUTHENTICATION,
)
from tornado import gen
import json
import terminado
import tornado.websocket
import re
Expand Down Expand Up @@ -78,8 +86,6 @@ def open(self, *args, **kwargs):
oauth_token.user
if valid:
user = oauth_token.user
else:
raise Exception('Invalid token')

self.term_name = term_name if term_name else 'tty'
if user:
Expand All @@ -91,6 +97,33 @@ def open(self, *args, **kwargs):
self.terminal.clients.append(self)
self.__initiate_terminal(self.terminal)

@gen.coroutine
def on_message(self, raw_message):
message = json.loads(raw_message)

api_key = message.get('api_key')
token = message.get('token')
command = message.get('command')

if REQUIRE_USER_AUTHENTICATION or is_disable_pipeline_edit_access():
valid = False

if api_key and token:
oauth_client = Oauth2Application.query.filter(
Oauth2Application.client_id == api_key,
).first()
if oauth_client:
oauth_token, valid = authenticate_client_and_token(oauth_client.id, token)
valid = valid and \
oauth_token and \
oauth_token.user and \
has_at_least_editor_role(oauth_token.user)
if not valid or is_disable_pipeline_edit_access():
return self.send_json_message(
['stdout', f'{command[1]}\nUnauthorized access to the terminal.'])

super().on_message(json.dumps(command))

def __initiate_terminal(self, terminal):
self.send_json_message(["setup", {}])
self._logger.info("TermSocket.open: Opened %s", self.term_name)
Expand Down

0 comments on commit f63cd00

Please sign in to comment.