extract code for getting federated users

core/extract/extractor.py

@@ -12,6 +12,7 @@
 import re
 import sys
 import traceback
+import pandas as pd
 
 import dateutil.parser
 import redshift_connector
@@ -81,6 +82,36 @@ def get_extract(self, log_location, start_time, end_time):
                 return self.cloudwatch_extractor.get_extract_from_cloudwatch(start_time, end_time)
             else:
                 return self.local_extractor.get_extract_locally(log_location, start_time, end_time)
+
+    def get_federated_objects(self):
+
+        region = self.config['region']
+        user = self.config['master_username']
+        database_name = self.config['source_cluster_endpoint'].split("/")[-1]
+        cluster_id = self.config['source_cluster_endpoint'].split(".")[0]
+        host = self.config['source_cluster_endpoint'].split(":")[0]
+        port = self.config['source_cluster_endpoint'].split(":")[1].split("/")[0]
+
+        df = pd.DataFrame(columns=['query tag','federated user','config','misc','entity','query/ACL'])
+
+        with open("/home/devsaba/code/redshift-test-drive/core/sql/federated_queries.sql",'r') as f:
+            queries_raw = f.read()
+        queries = queries_raw.split('\n\n\n')
+
+        creds = aws_service_helper.redshift_get_cluster_credentials(region=region, user=user, database_name=database_name,
+                                                                    cluster_id=cluster_id)
+
+        conn = util.db_connect(username=creds['DbUser'],password=creds['DbPassword'],
+                                            host=host,port=port, database=database_name)
+        cursor = conn.cursor()
+        for query in queries:
+            cursor.execute(query)
+            query_result = cursor.fetchall()
+
+            temp_df = pd.DataFrame.from_records(query_result,columns=['query tag','federated user','config','misc','entity','query/ACL'])
+            df = pd.concat([df,temp_df],axis=0,ignore_index=True)
+
+        return df            
 
     def save_logs(
         self,
@@ -215,6 +246,20 @@ def save_logs(
             replacements_file.write(replacements_string)
             replacements_file.close()
 
+        # Retrieve federated user objects
+        dataset = self.get_federated_objects()
+
+        if output_directory.startswith("s3://"):
+            dataset.to_csv("/tmp/federated-user-data.csv")
+            s3_object = output_directory[5:].partition("/")
+            bucket = s3_object[0]
+            prefix = s3_object[2]
+            object_name = '/federated-user-data.csv'
+
+            aws_service_helper.s3_upload('/tmp/federated-user-data.csv',bucket, f'{prefix}{object_name}' )
+        else:
+            dataset.to_csv(f'{output_directory}/federated-user-data.csv')
+
     def get_sql_connections_replacements(self, last_connections, log_items):
         # transactions has form { "xid": xxx, "pid": xxx, etc..., queries: [] }
         sql_json = {"transactions": OrderedDict()}

core/sql/federated_queries.sql

@@ -0,0 +1,175 @@
+/*Users*/
+select 'Users'::text as query_tag,
+  usename, 
+  null::text as col_placeholder1,
+  null::text as col_placeholder2,
+  null::text as col_placeholder3,
+  'create user ' || quote_ident('IAM' || regexp_substr(usename, '([^IAM:].*)')) || ' password disable;'
+FROM pg_user_info u
+where usename <> 'rdsdb'
+  and usesuper = 'f'
+  and usename like 'IAM:%';
+
+
+/*Add users and groups*/
+select 'Add users and groups'::text as query_tag, 
+       pg_get_userbyid(grolist[i]) as username,
+       groname,
+       null::text as col_placeholder1,
+       null::text as col_placeholder2,
+       'alter group ' || groname || ' add user ' ||
+       quote_ident('IAM' || regexp_substr(pg_get_userbyid(grolist[i]), '([^IAM:].*)')) || ';'       
+from (SELECT generate_series(1, array_upper(grolist, 1)) AS i, grolist, groname FROM pg_group) temp
+where grolist[i] not in (select usesysid from pg_user where usename = 'rdsdb' or usesuper = 't')
+  and username like 'IAM:%';
+
+
+/*Add users to roles */
+with data as (
+  select quote_ident('IAM' || regexp_substr(user_name, '([^IAM:].*)')) as username,
+                role_name  as rolename,
+                admin_option as adminOption,
+                'user'::text as type
+         from svv_user_grants
+         where user_name like 'IAM:%'
+)
+select 'Add users to roles'::text as query_tag,
+       username,
+       rolename,
+       adminOption,
+       null::text as col_placeholder1,
+       case
+           when type = 'user' and adminOption = 'false' then 'grant role ' || rolename || ' to ' || username || ';'
+           else 'grant role ' || rolename || ' to ' || username || ' with admin option;' end as query
+from data;
+
+
+/*User configuration*/
+select 'User configuration'::text as query_tag,
+       usename,
+       useconfig[i],
+       null::text as col_placeholder1,
+       null::text as col_placeholder2,
+       'alter user ' || quote_ident('IAM' || regexp_substr(usename, '([^IAM:].*)')) || ' set ' ||
+       case
+           when split_part(useconfig[i], '=', 1) = 'TimeZone' then split_part(useconfig[i], '=', 1) || '=''' ||
+                                                                   split_part(useconfig[i], '=', 2) || ''''
+           else useconfig[i] end || ';'
+from (SELECT generate_series(1, array_upper(useconfig, 1)) AS i, useconfig, usename
+      FROM pg_user
+      where usename like 'IAM:%'
+        -- and usesuper = 'f'
+     ) temp;
+
+
+/*User profile*/
+select 'User profiles'::text as query_tag,
+       usename,
+         null::text as col_placeholder1,
+         null::text as col_placeholder2,
+         null::text as col_placeholder3,
+       'alter user ' || quote_ident('IAM' || regexp_substr(usename, '([^IAM:].*)')) ||
+       decode(usecreatedb, 't', ' createdb', '') ||
+       decode(usesuper, 't', ' createuser', '') || nvl(' connection limit ' || useconnlimit, '') ||
+       nvl(' valid until ''' || valuntil || '''', '') || ';'
+from (SELECT u.usename,
+             u.usesysid,
+             u.usecreatedb,
+             u.usesuper,
+             u.useconnlimit,
+             decode(valuntil, 'infinity'::abstime, 'infinity', 'invalid'::abstime, null,
+                    to_char(valuntil, 'YYYY-MM-DD HH24:MI:SS')) AS valuntil
+      FROM pg_user_info u
+      where usename like 'IAM:%'
+     ) temp;
+
+
+/*User syslog access*/
+select 'syslogaccess'::text as query_tag,
+usename,
+null::text as col_placeholder1,
+null::text as col_placeholder2,
+null::text as col_placeholder3,
+'alter user ' || quote_ident('IAM:'|| regexp_substr(usename, '([^IAM:].*)')) || decode(syslogaccess,'UNRESTRICTED',' syslog access UNRESTRICTED') || ';' as query_statement
+from(
+select usename,
+        usesysid,
+        syslogaccess
+from svl_user_info
+where usename like 'IAM:%')
+where query_statement is not null;
+
+
+/*privileges on DB*/
+SELECT 'privileges on DB'::text as query_tag,
+       QUOTE_IDENT(pg_get_userbyid(b.datdba))::text AS objowner,
+       null::text                                   AS schemaname,
+       QUOTE_IDENT(b.datname)::text                 AS objname,
+       'database'::text                             AS objtype,
+       array_to_string(b.datacl, ',')::text         AS aclstring
+FROM pg_database b
+where pg_get_userbyid(b.datdba) like 'IAM:%'
+and aclstring is not null;
+
+
+/*privileges on schema*/
+SELECT 'privileges on schema'::text as query_tag,
+       QUOTE_IDENT(pg_get_userbyid(b.nspowner))::text AS objowner,
+       null::text                                     AS schemaname,
+       QUOTE_IDENT(b.nspname)::text                   AS objname,
+       'schema'::text                                 AS objtype,
+       array_to_string(b.nspacl, ',')::text           AS aclstring
+FROM pg_namespace b
+where QUOTE_IDENT(b.nspname) not ilike 'pg\_temp\_%' and
+pg_get_userbyid(b.nspowner) like 'IAM:%' 
+and aclstring is not null;
+
+
+/*privileges on tables*/
+SELECT 'privileges on tables'::text as query_tag,
+       QUOTE_IDENT(pg_get_userbyid(b.relowner))::text AS objowner,
+       QUOTE_IDENT(trim(c.nspname))::text             AS schemaname,
+       QUOTE_IDENT(b.relname)::text                   AS objname,
+       'table'::text                                  AS objtype,
+       array_to_string(b.relacl, ',')::text           AS aclstring
+FROM pg_class b
+         join pg_namespace c on b.relnamespace = c.oid
+  where pg_get_userbyid(b.relowner) like 'IAM:%'
+  and aclstring is not null;
+
+
+/*privileges on languages*/
+SELECT 'privileges on languages'::text as query_tag,
+       null::text                     AS objowner,
+       null::text                     AS schemaname,
+       lanname::text                  AS objname,
+       'language'::text               AS objtype,
+       array_to_string(b.lanacl, ',') AS aclstring
+FROM pg_language b
+where lanacl is not null;
+
+
+/*privileges on functions*/
+SELECT 'privileges on functions'::text as query_tag,
+       QUOTE_IDENT(pg_get_userbyid(b.proowner))::text                          AS objowner,
+       QUOTE_IDENT(trim(c.nspname))::text                                      AS schemaname,
+       QUOTE_IDENT(proname) || '(' || oidvectortypes(proargtypes) || ')'::text AS objname,
+       'function'::text                                                        AS objtype,
+       array_to_string(b.proacl, ',')                                          AS aclstring
+FROM pg_proc b
+         join pg_namespace c on b.pronamespace = c.oid
+where pg_get_userbyid(b.proowner) like 'IAM:%'
+and aclstring is not null;
+
+
+/*default ACL privileges*/
+SELECT 'default ACL privileges'::text as query_tag,
+       QUOTE_IDENT(pg_get_userbyid(b.defacluser))::text               AS objowner,
+       QUOTE_IDENT(trim(c.nspname))::text                             AS schemaname,
+       decode(b.defaclobjtype, 'r', 'tables', 'f', 'functions')::text AS objname,
+       'default acl'::text                                            AS objtype,
+       array_to_string(b.defaclacl, ',')::text                        AS aclstring
+FROM pg_default_acl b
+         left join pg_namespace c on b.defaclnamespace = c.oid
+where pg_get_userbyid(b.defacluser) like 'IAM:%'
+and aclstring is not null;