Class: Db

Inherits:
Object
  • Object
show all
Defined in:
lib/service/db.rb,
lib/domain/tag/db.rb,
lib/domain/site/db.rb,
lib/domain/asset/db.rb,
lib/domain/report/db.rb,
lib/domain/country/db.rb,
lib/domain/project/db.rb,
lib/domain/software/db.rb,
lib/domain/site_target/db.rb,
lib/domain/vulnerability/db.rb,
lib/domain/operating_system/db.rb,
lib/domain/shared_credential/db.rb

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(host:, db:, user:, port:) ⇒ Db

Returns a new instance of Db.



11
12
13
14
15
16
17
 
# File 'lib/service/db.rb', line 11

def initialize(host:, db:, user:, port:)
  @host = host
  @db = db
  @user = user
  @port = port
  @connection = PG.connect(host:, dbname: db, port:, user:)
end

Instance Attribute Details

#connection ⇒ Object (readonly)

Returns the value of attribute connection.



9
10
11
 
# File 'lib/service/db.rb', line 9

def connection
  @connection
end

#db ⇒ Object (readonly)

Returns the value of attribute db.



9
10
11
 
# File 'lib/service/db.rb', line 9

def db
  @db
end

#host ⇒ Object (readonly)

Returns the value of attribute host.



9
10
11
 
# File 'lib/service/db.rb', line 9

def host
  @host
end

#port ⇒ Object (readonly)

Returns the value of attribute port.



9
10
11
 
# File 'lib/service/db.rb', line 9

def port
  @port
end

#user ⇒ Object (readonly)

Returns the value of attribute user.



9
10
11
 
# File 'lib/service/db.rb', line 9

def user
  @user
end

Instance Method Details

#all(model_class) ⇒ Object 



19
20
21
22
23
24
25
 
# File 'lib/service/db.rb', line 19

def all(model_class)
  sql = "SELECT * FROM #{model_class.table_name}"
  result = @connection.exec(sql)
  result.map do |row|
    model_class.new(row)
  end
end

#bulk_copy(models = []) ⇒ Object 



84
85
86
87
88
89
90
91
92
 
# File 'lib/service/db.rb', line 84

def bulk_copy(models = [])
  return if models.empty?

  model = models.first
  table_name = model.class.table_name
  columns = model.class.columns
  data = models.map(&:to_csv)
  bulk_copy_csv(table_name:, columns:, data:)
end

#bulk_copy_csv(table_name:, columns: [], data: []) ⇒ Object 



94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
 
# File 'lib/service/db.rb', line 94

def bulk_copy_csv(table_name:, columns: [], data: [])
  column_names = columns.empty? ? '' : "( #{columns.join(',')})"
  progress_bar = ProgressBar.create(
    title: "Copy to #{table_name}",
    total: data.length,
    format: '%a %B %p%% %t'
  )
  @connection.copy_data("COPY #{table_name} #{column_names} from STDIN with (FORMAT csv)") do
    data.each do |row|
      sanitized_row = row.map { |cell| escape_csv_value(cell) }
      connection.put_copy_data(sanitized_row.join(',') + "\n")
      progress_bar.increment
    end
  end
end

#countries ⇒ Object 



10
11
12
 
# File 'lib/domain/country/db.rb', line 10

def countries
  @countries ||= fetch_all_countries
end

#fetch_cmdb_assets ⇒ Object 



6
7
8
9
10
 
# File 'lib/domain/tag/db.rb', line 6

def fetch_cmdb_assets
  fetch_view('utr_asset_view') do |row|
    CmdbAsset.from_csv(row)
  end
end

#fetch_cyberark_shared_credentials(_country_code) ⇒ Object 



13
14
15
16
17
 
# File 'lib/domain/site_target/db.rb', line 13

def fetch_cyberark_shared_credentials(_country_code)
  credentials = fetch_shared_credentials
  credentials.select do |credential|
  end
end

#fetch_out_of_vm_scope_asset_ids ⇒ Object 



12
13
14
15
16
17
18
 
# File 'lib/domain/asset/db.rb', line 12

def fetch_out_of_vm_scope_asset_ids
  ids = []
  fetch_view('asset_out_of_vm_scope_view') do |row|
    ids << row[:id].to_i
  end
  ids
end

#fetch_shared_credentials ⇒ Object 



7
8
9
10
11
 
# File 'lib/domain/site_target/db.rb', line 7

def fetch_shared_credentials
  fetch_view('country_view') do |row|
    Country.from_csv(row)
  end
end

#fetch_utr_site_from_cmdb ⇒ Object 



7
8
9
10
11
 
# File 'lib/domain/site/db.rb', line 7

def fetch_utr_site_from_cmdb
  fetch_view('utr_site_view') do |row|
    CmdbSite.from_csv(row)
  end
end

#fetch_view(view, &block) ⇒ Object 



153
154
155
156
157
158
159
 
# File 'lib/service/db.rb', line 153

def fetch_view(view, &block)
  ENV['PGPASSFILE'] = '/Users/ckyony/.pgpass'
  command = "psql -d #{db} -U #{user} -h '#{host}' -p #{port} -c \"\\copy (SELECT * FROM #{view}) to STDOUT with CSV HEADER;\""
  result = `#{command}`
  rows = CSV.new(result, headers: true, header_converters: :symbol)
  rows.map(&block)
end

#find_by(model_class, attribute, value) ⇒ Object 



35
36
37
38
39
40
41
42
43
44
 
# File 'lib/service/db.rb', line 35

def find_by(model_class, attribute, value)
  sql = "SELECT * FROM #{model_class.table_name} WHERE #{attribute} = $1"
  result = @connection.exec_params(sql, [value])
  result.map do |row|
    model_class.new(row)
  end
  return nil if result.ntuples.zero?

  model_class.new(result.first.symbolize_keys)
end

#grant_existing_view_access(user, role, schema) ⇒ Object 



53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
 
# File 'lib/service/db.rb', line 53

def grant_existing_view_access(user, role, schema)
  query = <<-SQL
  DO $$
  DECLARE
    r RECORD;
  BEGIN
    FOR r IN (SELECT table_schema, table_name
              FROM information_schema.views v
              JOIN pg_views pv
                ON v.table_schema = pv.schemaname
                AND v.table_name = pv.viewname
              WHERE table_schema = '#{schema}'
                AND pv.viewowner = '#{role}'
              ) LOOP
      EXECUTE 'GRANT SELECT ON ' || quote_ident(r.table_schema) || '.' || quote_ident(r.table_name) || ' TO ' || '#{user}';
    END LOOP;
  END $$ LANGUAGE plpgsql;
  SQL

  puts "query #{query}"

  @connection.exec(query)
end

#grant_future_view_access(user, role, schema) ⇒ Object 



77
78
79
80
81
82
 
# File 'lib/service/db.rb', line 77

def grant_future_view_access(user, role, schema)
  query = <<-SQL
  ALTER DEFAULT PRIVILEGES FOR ROLE #{role} IN SCHEMA #{schema} GRANT SELECT ON TABLES TO #{user};
  SQL
  @connection.exec(query)
end

#grant_view_access(user, role, schema = 'public') ⇒ Object 



46
47
48
49
50
51
 
# File 'lib/service/db.rb', line 46

def grant_view_access(user, role, schema = 'public')
  puts "user, role,schema #{[user, role, schema]}"
  grant_existing_view_access(user, role, schema)
  grant_future_view_access(user, role, schema)
  @connection.close
end

#operating_system_found_on_assets ⇒ Object 



8
9
10
 
# File 'lib/domain/operating_system/db.rb', line 8

def operating_system_found_on_assets
  select(OperatingSystem, 'operating_system_view')
end

#product_in_cmdb_eos ⇒ Object 



23
24
25
 
# File 'lib/domain/software/db.rb', line 23

def product_in_cmdb_eos
  all(CmdbEos)
end

#save_cisa_kev(vulnerability) ⇒ Object 



7
8
9
 
# File 'lib/domain/vulnerability/db.rb', line 7

def save_cisa_kev(vulnerability)
  upsert(vulnerability)
end

#save_microsoft_product_lifecycle(models) ⇒ Object 



27
28
29
 
# File 'lib/domain/software/db.rb', line 27

def save_microsoft_product_lifecycle(models)
  bulk_copy(models)
end

#save_project_report_details(details = []) ⇒ Object 



6
7
8
 
# File 'lib/domain/project/db.rb', line 6

def save_project_report_details(details = [])
  bulk_copy(details)
end

#select(model_class, view_name) ⇒ Object 



27
28
29
30
31
32
33
 
# File 'lib/service/db.rb', line 27

def select(model_class, view_name)
  sql = "SELECT * FROM #{view_name}"
  result = @connection.exec(sql)
  result.map do |row|
    model_class.new(row)
  end
end

#software_database ⇒ Object 



21
 
# File 'lib/domain/software/db.rb', line 21

def software_database; end

#software_found_on_assets ⇒ Object 



8
9
10
 
# File 'lib/domain/software/db.rb', line 8

def software_found_on_assets
  select(Software, 'software_view')
end

#softwares(type) ⇒ Object 



12
13
14
15
16
17
18
19
 
# File 'lib/domain/software/db.rb', line 12

def softwares(type)
  list = select(Software, 'software_view')
  return list if type == 'all'

  return list.select(&:db?) if type == 'db'

  list.reject(&:db?)
end

#upsert(model) ⇒ Object 



110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
 
# File 'lib/service/db.rb', line 110

def upsert(model)
  table_name = model.class.table_name
  columns = model.class.columns
  column_names = columns.join(', ')
  placeholders = columns.map.with_index(1) { |_, i| "$#{i}" }.join(', ')
  conflict_target = model.class.primary_key.join(', ') # Unique identifier assumption
  # puts "---- conflict_target #{conflict_target}"
  # puts "---- table_name: #{table_name}"
  # puts "---- model.class: #{model.class}"
  # puts "---- model.class.primary_key: #{model.class.primary_key}"

  update_assignments = columns.map do |column|
    "#{column} = EXCLUDED.#{column}"
  end.join(', ')

  sql = <<-SQL
    INSERT INTO #{table_name} (#{column_names})
    VALUES (#{placeholders})
    ON CONFLICT (#{conflict_target}) DO UPDATE SET
    #{update_assignments};
  SQL

  @connection.exec_params(sql, model.to_csv)
rescue PG::Error => e
  puts "An error occurred: #{e.message}"
end

#upsert_without_primary_key(model) ⇒ Object 



137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
 
# File 'lib/service/db.rb', line 137

def upsert_without_primary_key(model)
  table_name = model.class.table_name
  columns = model.class.columns
  column_names = columns.join(', ')
  placeholders = columns.map.with_index(1) { |_, i| "$#{i}" }.join(', ')

  sql = <<-SQL
    INSERT INTO #{table_name} (#{column_names})
    VALUES (#{placeholders})
  SQL

  @connection.exec_params(sql, model.to_csv)
rescue PG::Error => e
  puts "An error occurred: #{e.message}"
end