Coverage for mindsdb / utilities / auth.py: 0%

1import requests 

2import textwrap 

3 

4from mindsdb.utilities.config import Config 

5 

6 

7def get_aws_meta_data() -> dict: 

8 ''' returns aws metadata for current instance 

9 

10 Returns: 

11 dict: aws metadata 

12 ''' 

13 aws_meta_data = { 

14 'public-hostname': None, 

15 'ami-id': None, 

16 'instance-id': None 

17 } 

18 aws_token = requests.put("http://169.254.169.254/latest/api/token", headers={'X-aws-ec2-metadata-token-ttl-seconds': '30'}).text 

19 for key in aws_meta_data.keys(): 

20 resp = requests.get( 

21 f'http://169.254.169.254/latest/meta-data/{key}', 

22 headers={'X-aws-ec2-metadata-token': aws_token}, 

23 timeout=1 

24 ) 

25 if resp.status_code != 200: 

26 continue 

27 aws_meta_data[key] = resp.text 

28 if aws_meta_data['instance-id'] is None: 

29 raise Exception('That is not an AWS environment') 

30 return aws_meta_data 

31 

32 

33def register_oauth_client(): 

34 ''' register new oauth client if it is not existed 

35 ''' 

36 config = Config() 

37 aws_meta_data = get_aws_meta_data() 

38 

39 current_aws_meta_data = config.get('aws_meta_data', {}) 

40 oauth_meta = config.get('auth', {}).get('oauth') 

41 if oauth_meta is None: 

42 return 

43 

44 public_hostname = aws_meta_data['public-hostname'] 

45 if ( 

46 current_aws_meta_data.get('public-hostname') != public_hostname 

47 or oauth_meta.get('client_id') is None 

48 ): 

49 resp = requests.post( 

50 f'https://{oauth_meta["server_host"]}/auth/register_client', 

51 json={ 

52 'client_name': f'aws_marketplace_{public_hostname}', 

53 'client_uri': public_hostname, 

54 'grant_types': 'authorization_code', 

55 'redirect_uris': textwrap.dedent(f''' 

56 https://{public_hostname}/api/auth/callback 

57 https://{public_hostname}/api/auth/callback/cloud_home 

58 '''), 

59 'response_types': 'code', 

60 'scope': 'openid profile aws_marketplace', 

61 'token_endpoint_auth_method': 'client_secret_basic' 

62 }, 

63 timeout=10 

64 ) 

65 

66 if resp.status_code != 200: 

67 raise Exception(f'Wrong answer from auth server: {resp.status_code}, {resp.text}') 

68 keys = resp.json() 

69 Config().update({ 

70 'aws_meta_data': aws_meta_data, 

71 'auth': { 

72 'oauth': { 

73 'client_id': keys['client_id'], 

74 'client_secret': keys['client_secret'] 

75 } 

76 } 

77 })