async def create_staging_table(conn: Connection):
"""Create a staging table for the cards, empties it if it exists.
Note that the table is unlogged, so it's supposed to be used for this
operation only, then lost.
"""
async with conn.transaction():
await conn.execute("DROP TABLE IF EXISTS card_stg")
async with conn.transaction():
await conn.execute("""
CREATE UNLOGGED TABLE card_stg (
from_lang SMALLINT NOT NULL,
to_lang SMALLINT NOT NULL,
from_id INTEGER NOT NULL,
to_id INTEGER NOT NULL,
from_txt TEXT NOT NULL,
original_txt TEXT NOT NULL,
to_tokens TEXT[] NOT NULL
) PARTITION BY HASH (from_id, to_id);
""")
for i in range(10):
await conn.execute(f"""
CREATE UNLOGGED TABLE card_stg_h{i}
PARTITION OF card_stg
FOR VALUES WITH (MODULUS 10, REMAINDER {i});
""")
async def update(self,
con: Connection = None,
table: str = '',
where_dict: dict = None,
**update_params: dict) -> list:
values = ','.join(
["{}='{}'".format(k, v) for k, v in update_params.items()])
where = ''
if where_dict is not None:
where = ' where '
where += ' and '.join([
self.WHERE.format(key=k, value=v)
for k, v in where_dict.items()
])
query = self.UPDATE.format(table=table, values=values, where=where)
if not con:
pool = await self.get_pool()
async with pool.acquire() as con:
async with con.transaction():
results = await con.fetch(query)
else:
async with con.transaction():
results = await con.fetch(query)
return results
async def execute(self, con: Connection = None, query: str = ''):
if not con:
pool = await self.get_pool()
async with pool.acquire() as con:
async with con.transaction():
await con.execute(query)
else:
async with con.transaction():
await con.execute(query)
async def merge_tables(conn: Connection, p_id: int):
"""Merge the staging table and the final one."""
async with conn.transaction():
logger.debug('Detecting cards no longer valid...')
res = await conn.fetchrow(f"""
SELECT count(1) AS to_delete
FROM card_h{p_id} c
LEFT JOIN card_stg_h{p_id} s
USING (from_id, to_id)
WHERE s.from_id IS NULL
""")
if res['to_delete'] > 500:
# around 10 cards are deleted per day, too many may be an issue with the
# file, so better stop rather than deleting everything
raise ValueError(f"Suspect number of cards to delete: {res['to_delete']}")
if res['to_delete'] > 0:
logger.info(f"Found {res['to_delete']} invalid cards, deleting them")
await conn.execute(f"""
DELETE
FROM card_h{p_id}
WHERE
(from_lang, to_lang) IN
(
SELECT
c.from_lang,
c.to_lang
FROM card_h{p_id} c
LEFT JOIN card_stg_h{p_id} s
USING (from_id, to_id)
WHERE
s.from_id IS NULL)
""")
else:
logger.info('There were no cards to delete')
logger.info('Updating the cards which changed...')
res = await conn.fetchrow(f"""
UPDATE card_h{p_id} c
SET
from_txt = s.from_txt,
original_txt = s.original_txt,
to_tokens = s.to_tokens
FROM card_stg_h{p_id} s
WHERE
c.from_id = s.from_id
AND c.to_id = s.to_id
AND (c.from_txt, c.to_tokens) <> (s.from_txt, s.to_tokens)
""")
logger.info(f'Update successful: changes {res}')
logger.info('Inserting new cards...')
res = await conn.fetchrow(f"""
INSERT INTO card_h{p_id}
SELECT
s.*
FROM
card_stg_h{p_id} s
LEFT JOIN card_h{p_id} c
USING(from_id, to_id)
WHERE
c.from_lang IS NULL
""")
async def insertTop10Companies(connection: Connection):
async with connection.transaction():
await connection.execute('''
DELETE FROM empresas
''')
await connection.copy_records_to_table(
"empresas",
records=[(1, "Oi", "OIBR-C", "Rio de Janeiro, Rio de Janeiro",
"Telecomunicações", 10),
(2, "Braskem", "BRKM5.SAO", "São Paulo, São Paulo",
"Petroquímica", 9),
(3, "Itaúsa", "ITSA4.SAO", "Brasil",
"Finanças e indústria", 8),
(4, "JBS", "JBSS3.SAO", "São Paulo, São Paulo",
"Alimentício", 7),
(5, "Eletrobras", "ELET6.SAO",
"Rio de Janeiro, Rio de Janeiro", "Energia elétrica", 6),
(6, "Banco do Brasil", "BBAS3.SAO",
"Brasília, Distrito Federal", "Bancário", 5),
(7, "Vale", "VALE3.SAO", "Rio de Janeiro, Rio de Janeiro",
"Mineração", 4),
(8, "Banco Bradesco", "BBDC4.SAO", "Osasco, São Paulo",
"Bancário", 3),
(9, "Itaú Unibanco", "ITUB4.SAO", "São Paulo, São Paulo",
"Bancário", 2),
(10, "Petrobras", "PETR4.SAO",
"Rio de Janeiro, Rio de Janeiro", "Petróleo e gás", 1)])
async def db_import_rule(rule_id: str,
rule_cves: list,
conn: Connection,
rule_only=False):
"""Import single error key into database"""
try:
async with conn.transaction():
inserted = await conn.fetchrow(
"""INSERT INTO insights_rule (name, rule_only)
VALUES($1, $2)
ON CONFLICT (name) DO UPDATE
SET name = EXCLUDED.name, rule_only = EXCLUDED.rule_only
RETURNING id AS inserted""",
rule_id, rule_only)
RULES_CACHE[rule_id] = inserted[0]
to_insert = []
for rule_cve in rule_cves:
if rule_cve not in CVES_CACHE:
await db_import_cve(rule_cve, conn)
to_insert.append((inserted[0], CVES_CACHE[rule_cve]))
await conn.executemany(
"""INSERT INTO cve_rule_mapping (rule_id, cve_id)
VALUES ($1, $2) ON CONFLICT DO NOTHING""",
to_insert)
# pylint: disable=broad-except
except Exception:
DATABASE_ERROR.inc()
LOGGER.exception("Error during inserting rule: ")
async def delete(self,
con: Connection = None,
table: str = '',
where_dict: dict = None):
where = ' where '
where += ' and '.join(
[self.WHERE.format(key=k, value=v) for k, v in where_dict.items()])
query = self.DELETE.format(table=table, where=where)
if not con:
pool = await self.get_pool()
async with pool.acquire() as con:
async with con.transaction():
await con.execute(query)
else:
async with con.transaction():
await con.execute(query)
async def createCompanyTable(connection: Connection):
async with connection.transaction():
await connection.execute('''
CREATE TABLE IF NOT EXISTS empresas(
id bigserial NOT NULL,
nome character varying(50) NOT NULL,
simbolo_empresa character varying(20) PRIMARY KEY NOT NULL UNIQUE,
regiao character varying(50),
setor character varying(20),
rank int
)''')
async def createStockTable(connection: Connection):
async with connection.transaction():
await connection.execute('''
CREATE TABLE IF NOT EXISTS cotacoes(
preco float(3),
volume bigint,
data timestamp without time zone,
variacao float(3),
variacao_por_cento float(5),
simbolo_empresa character varying(20) NOT NULL,
CONSTRAINT cotacao_fk FOREIGN KEY(simbolo_empresa) REFERENCES empresas(simbolo_empresa) ON DELETE CASCADE
)''')
async def insert(self,
con: Connection = None,
table: str = '',
value_dict: dict = None):
columns = ",".join(value_dict.keys())
placeholder = ",".join(
['${}'.format(i) for i in range(1,
len(value_dict) + 1)])
query = self.INSERT.format(table=table,
columns=columns,
values=placeholder)
if not con:
pool = await self.get_pool()
async with pool.acquire() as con:
async with con.transaction():
result = await con.fetchrow(query, *value_dict.values())
else:
async with con.transaction():
result = await con.fetchrow(query, *value_dict.values())
return result
async def ingest_cards_file(
conn: Connection, jsonl_file: str, languages_ids: Dict[str, int]):
"""Insert a cards file in the staging table card_stg."""
pending = []
last_update = time()
for idx, line in enumerate(open(jsonl_file)):
card = json.loads(line)
try:
pending.append((
languages_ids[card['from_lang']],
languages_ids[card['to_lang']],
card['from_id'],
card['to_id'],
card['from_txt'],
card['original_txt'],
card['resulting_tokens'],
))
except KeyError as ke:
# missing language or weird line
logger.warning(f'Error processing a row: {line}: {ke}')
if len(pending) > 2000:
async with conn.transaction():
await conn.copy_records_to_table(
'card_stg',
records=pending,
)
pending = []
if time() > last_update + 60:
logger.debug(f'Ingested {idx} so far...')
last_update = time()
logger.debug('Almost done, writing the remaining entries...')
async with conn.transaction():
await conn.copy_records_to_table(
'card_stg',
records=pending,
)
def _get_query(self, connection: Connection,
params: Dict[str, str]) -> Awaitable[Record]:
"""Internal method. Return query to fetch."""
query: Dict[str, List[str]] = {
'statements': [],
'params': [],
}
self._prepare_params(query, params)
return connection.fetch(
'SELECT * FROM pools' +
(' WHERE ' +
', '.join(query['statements']) if query['statements'] else ''),
*query['params'])
async def _execute_file(migration_file: MigrationFile, conn: Connection):
log.info(f"EXECUTE {migration_file.stamp} {migration_file.name}")
try:
async with conn.transaction():
with open(migration_file.path) as stream:
migration_sql = stream.read()
await conn.execute(migration_sql)
insert_sql = "INSERT INTO arctic_tern_migrations VALUES ($1, $2, $3, now())"
await conn.execute(insert_sql, migration_file.stamp,
migration_file.name, migration_file.hash_)
except PostgresError as e:
log.error(e)
raise e
async def delete_staging_table(conn: Connection):
"""Delete the staging table."""
async with conn.transaction():
await conn.execute("DROP TABLE card_stg")
async def store_language_codes(conn: Connection) -> Dict[str, int]:
"""Update the language codes and return the code -> id dictionary.
The DB is updated if the name changed or a language is new.
If a language is in the DB but not in the hardcoded list, an exception is raised.
Returns
-------
Dict[str, int]
A dictionary mapping the ISO code of a language with its ID
Raises
------
ValueError
In case a language was in the DB but not in the hardcoded list
"""
languages_ids = {} # ISO -> (id, name)
max_id = 0
async with conn.transaction():
res = await conn.fetch('SELECT id, name, iso693_3 FROM language')
for lng in res:
languages_ids[lng['iso693_3']] = (lng['id'], lng['name'])
if max_id < lng['id']:
max_id = lng['id']
for iso, name in ISO_693_3.items():
if iso in languages_ids:
pre_name = languages_ids[iso][1]
# different name? update it
if pre_name != name:
logger.info(
f'Language with ISO {iso} is now named {name} instead of {pre_name}'
)
await conn.execute(
"""UPDATE language
SET name = $1
WHERE iso693_3 = $2
""",
name,
iso,
)
else:
# not there, insert it
max_id += 1
logger.info(f'New language, {iso} => {name} will have id {max_id}')
languages_ids[iso] = (max_id, name)
await conn.execute(
"""INSERT INTO language(
id,
iso693_3,
name)
VALUES ($1, $2, $3)""",
max_id, iso, name
)
# now all of ISO_693_3 elements are in the DB and in language_ids
# let's check for language_ids which are gone
gone_languages = set(languages_ids.keys()).difference(ISO_693_3.keys())
if len(gone_languages) > 0:
raise ValueError(
f'Some languages are in the DB but are unknow! They are {gone_languages}'
)
return {iso: id_ for iso, (id_, _) in languages_ids.items()}
def notified(cnx: Connection, pid: int, channel: str, payload: str):
asyncio.gather(cnx.close(), q.put(channel))
async def db_import_cve(cve: str, conn: Connection):
"""Import missing CVE metadata into database"""
# Attempt to get CVE data from VMAAS
LOGGER.info("Attempting to get %s data from VMAAS", cve)
cve_request_endpoint = format_vmaas_cve_endpoint(cve)
cve_page = await vmaas_request(cve_request_endpoint, method="GET")
success = True
try:
cve_data = cve_page["cve_list"][cve]
except (KeyError, TypeError):
LOGGER.error("Error importing %s from VMAAS", cve)
success = False
try:
async with conn.transaction():
if success:
impact_id_map = {}
async for row in conn.cursor(
"SELECT name, id FROM cve_impact"):
impact_name, impact_id = row
impact_id_map[impact_name] = impact_id
cve_row = construct_cve_row(cve_data,
impact_id_map,
asyncpg_type=True)
inserted = await conn.fetchrow(
"""INSERT INTO cve_metadata
(cve, description, impact_id, public_date, modified_date,
cvss3_score, cvss3_metrics, cvss2_score, cvss2_metrics, redhat_url,
secondary_url, advisories_list)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
ON CONFLICT (cve) DO UPDATE
SET
description = EXCLUDED.description,
impact_id = EXCLUDED.impact_id,
public_date = EXCLUDED.public_date,
modified_date = EXCLUDED.modified_date,
cvss3_score = EXCLUDED.cvss3_score,
cvss3_metrics = EXCLUDED.cvss3_metrics,
cvss2_score = EXCLUDED.cvss2_score,
cvss2_metrics = EXCLUDED.cvss2_metrics,
redhat_url = EXCLUDED.redhat_url,
secondary_url = EXCLUDED.secondary_url,
advisories_list = EXCLUDED.advisories_list
RETURNING id AS inserted""", *cve_row)
# if vmaas was not available insert empty cve
else:
inserted = await conn.fetchrow(
"""INSERT INTO cve_metadata (cve, description, impact_id)
VALUES ($1, $2, $3)
ON CONFLICT (cve) DO UPDATE
SET cve = $1
RETURNING id AS inserted""",
cve, 'unknown', 0)
CVES_CACHE[cve] = inserted[0]
# pylint: disable=broad-except
except Exception:
DATABASE_ERROR.inc()
LOGGER.exception("Error during inserting CVE: ")
async def db_import_rule_hits(conn: Connection, rh_account_id: int,
inventory_id: str, system_id: int,
rule_hits: dict) -> None:
"""Associate rules hits with system"""
# pylint: disable=too-many-branches, too-many-statements
active_rules = set()
rules_playbook_count = {}
async for row in conn.cursor(
"""SELECT ir.id, ir.playbook_count, ir.active FROM insights_rule ir"""
):
rule_id, rule_playbook_count, rule_active = row
rules_playbook_count[rule_id] = rule_playbook_count
if rule_active:
active_rules.add(rule_id)
to_update = []
rule_hits_cves_ids = tuple(rule_hits.keys())
system_cves = []
async for row in conn.cursor(
"""SELECT sv.id, sv.cve_id, sv.rule_id, sv.when_mitigated, sv.advisory_available, cm.cve
FROM system_vulnerabilities sv
JOIN cve_metadata cm ON sv.cve_id = cm.id
WHERE system_id = $1 AND rh_account_id = $2""",
system_id, rh_account_id):
row_id, cve_id, rule_id, when_mitigated, advisory_available, cve = row
playbook_count = rules_playbook_count.get(
rule_hits.get(cve_id, {}).get('id'))
mitigation_reason = rule_hits.get(cve_id, {}).get('mitigation_reason')
remediation_type_id = get_available_remediation_type(
advisory_available, when_mitigated, mitigation_reason,
playbook_count)
if cve_id in rule_hits_cves_ids:
if 'mitigation_reason' in rule_hits[cve_id]:
if rule_hits[cve_id][
'id'] not in active_rules and not when_mitigated:
system_cves.append(cve)
to_update.append(
(row_id, rh_account_id, rule_hits[cve_id]['id'], None,
rule_hits[cve_id]['mitigation_reason'],
advisory_available, remediation_type_id))
else:
to_update.append(
(row_id, rh_account_id, rule_hits[cve_id]['id'],
rule_hits[cve_id]['details'], None, advisory_available,
remediation_type_id))
if rule_hits[cve_id][
'id'] in active_rules or not when_mitigated:
system_cves.append(cve)
del rule_hits[
cve_id] # rule hits become dict of completely new system cves
elif rule_id:
to_update.append((row_id, rh_account_id, None, None, None,
advisory_available, remediation_type_id))
if not when_mitigated:
system_cves.append(cve)
elif not when_mitigated:
system_cves.append(cve)
if rule_hits:
await insert_new_cves(conn, rh_account_id, system_id, rule_hits,
rules_playbook_count)
system_cves.extend([
rule_hits[cve]['cve_name'] for cve in rule_hits
if rule_hits[cve]['id'] in active_rules
and 'details' in rule_hits[cve]
])
if to_update:
await update_cves(conn, to_update, rh_account_id)
await _update_system(system_id, len(system_cves), conn)
send_remediations_update(REMEDIATIONS_PRODUCER, inventory_id, system_cves)