def WriteUserNotification(self, notification, cursor=None): """Writes a notification for a given user.""" # Copy the notification to ensure we don't modify the source object. notification = notification.Copy() if not notification.timestamp: notification.timestamp = rdfvalue.RDFDatetime.Now() args = { "username_hash": mysql_utils.Hash(notification.username), "timestamp": mysql_utils.RDFDatetimeToMysqlString(notification.timestamp), "notification_state": int(notification.state), "notification": notification.SerializeToString(), } query = "INSERT INTO user_notification {columns} VALUES {values}".format( columns=mysql_utils.Columns(args), values=mysql_utils.NamedPlaceholders(args)) try: cursor.execute(query, args) except MySQLdb.IntegrityError: raise db.UnknownGRRUserError(notification.username)
def WriteGRRUser(self, username, password=None, ui_mode=None, canary_mode=None, user_type=None, cursor=None): """Writes user object for a user with a given name.""" values = { "username": username, "username_hash": mysql_utils.Hash(username) } if password is not None: values["password"] = password.SerializeToString() if ui_mode is not None: values["ui_mode"] = int(ui_mode) if canary_mode is not None: # TODO(amoser): This int conversion is dirty but necessary with # the current MySQL driver. # TODO: We can remove this once the bug is fixed. values["canary_mode"] = int(bool(canary_mode)) if user_type is not None: values["user_type"] = int(user_type) query = "INSERT INTO grr_users {cols} VALUES {vals}".format( cols=mysql_utils.Columns(values), vals=mysql_utils.NamedPlaceholders(values)) updates = ", ".join("{0} = VALUES({0})".format(col) for col in values) query += " ON DUPLICATE KEY UPDATE " + updates cursor.execute(query, values)
def testDictUsesKeys(self): self.assertIn( mysql_utils.Columns({ "bar": 42, "baz": 42, "foo": 42 }), ["(`bar`, `baz`, `foo`)"])
def _Insert(cursor, table, values): """Inserts one or multiple rows into the given table. Args: cursor: The MySQL cursor to perform the insertion. table: The table name, where rows should be inserted. values: A list of dicts, associating column names to values. """ precondition.AssertIterableType(values, dict) if not values: # Nothing can be INSERTed with empty `values` list. return column_names = list(sorted(values[0])) for value_dict in values: if set(column_names) != set(value_dict): raise ValueError("Given value dictionaries must have identical keys. " "Expecting columns {!r}, but got value {!r}".format( column_names, value_dict)) query = "INSERT IGNORE INTO %s {cols} VALUES {vals}" % table query = query.format( cols=mysql_utils.Columns(column_names), vals=mysql_utils.Placeholders(num=len(column_names), values=len(values))) values_list = [] for values_dict in values: values_list.extend(values_dict[column] for column in column_names) cursor.execute(query, values_list)
def WriteSignedBinaryReferences(self, binary_id, references, cursor=None): """Writes blob references for a signed binary to the DB.""" args = { "binary_type": binary_id.binary_type.SerializeToDataStore(), "binary_path": binary_id.path, "binary_path_hash": mysql_utils.Hash(binary_id.path), "timestamp": mysql_utils.RDFDatetimeToMysqlString(rdfvalue.RDFDatetime.Now()), "blob_references": references.SerializeToString() } query = """ INSERT INTO signed_binary_references {cols} VALUES {vals} ON DUPLICATE KEY UPDATE timestamp = VALUES(timestamp), blob_references = VALUES(blob_references) """.format( cols=mysql_utils.Columns(args), vals=mysql_utils.NamedPlaceholders(args)) cursor.execute(query, args)
def _GrantApproval(self, requestor_username, approval_id, grantor_username, cursor): """Grants approval for a given request.""" grant_args = { "username_hash": mysql_utils.Hash(requestor_username), "approval_id": approval_id, "grantor_username_hash": mysql_utils.Hash(grantor_username), } grant_query = ( "INSERT INTO approval_grant {columns} VALUES {values}".format( columns=mysql_utils.Columns(grant_args), values=mysql_utils.NamedPlaceholders(grant_args))) cursor.execute(grant_query, grant_args)
def WriteSignedBinaryReferences(self, binary_id, references, cursor=None): """Writes blob references for a signed binary to the DB.""" args = { "binary_type": binary_id.binary_type.SerializeToWireFormat(), "binary_path": binary_id.path, "binary_path_hash": mysql_utils.Hash(binary_id.path), "blob_references": references.SerializeToBytes() } query = """ INSERT INTO signed_binary_references {cols} VALUES {vals} ON DUPLICATE KEY UPDATE blob_references = VALUES(blob_references) """.format(cols=mysql_utils.Columns(args), vals=mysql_utils.NamedPlaceholders(args)) cursor.execute(query, args)
def WriteAPIAuditEntry(self, entry, cursor=None): """Writes an audit entry to the database.""" args = { "username_hash": mysql_utils.Hash(entry.username), "router_method_name": entry.router_method_name, "details": entry.SerializeToString(), "timestamp": mysql_utils.RDFDatetimeToMysqlString(rdfvalue.RDFDatetime.Now()) } query = "INSERT INTO api_audit_entry {columns} VALUES {values}".format( columns=mysql_utils.Columns(args), values=mysql_utils.NamedPlaceholders(args)) cursor.execute(query, args)
def WriteApprovalRequest(self, approval_request, cursor=None): """Writes an approval request object.""" # Copy the approval_request to ensure we don't modify the source object. approval_request = approval_request.Copy() # Generate random approval id. approval_id_int = random.UInt64() now_str = mysql_utils.RDFDatetimeToMysqlString( rdfvalue.RDFDatetime.Now()) grants = approval_request.grants approval_request.grants = None args = { "username_hash": mysql_utils.Hash(approval_request.requestor_username), "approval_type": int(approval_request.approval_type), "subject_id": approval_request.subject_id, "approval_id": approval_id_int, "timestamp": now_str, "expiration_time": mysql_utils.RDFDatetimeToMysqlString( approval_request.expiration_time), "approval_request": approval_request.SerializeToString() } query = ( "INSERT INTO approval_request {columns} VALUES {values}".format( columns=mysql_utils.Columns(args), values=mysql_utils.NamedPlaceholders(args))) cursor.execute(query, args) for grant in grants: self._GrantApproval(approval_request.requestor_username, approval_id_int, grant.grantor_username, now_str, cursor) return _IntToApprovalID(approval_id_int)
def WriteClientGraphSeries( self, graph_series, client_label, timestamp, cursor=None, ): """Writes the provided graphs to the DB with the given client label.""" args = { "client_label": client_label, "report_type": graph_series.report_type.SerializeToDataStore(), "timestamp": mysql_utils.RDFDatetimeToMysqlString(timestamp), "graph_series": graph_series.SerializeToString(), } query = """ INSERT INTO client_report_graphs {cols} VALUES {vals} ON DUPLICATE KEY UPDATE graph_series = VALUES(graph_series) """.format(cols=mysql_utils.Columns(args), vals=mysql_utils.NamedPlaceholders(args)) cursor.execute(query, args)
def testMany(self): self.assertEqual(mysql_utils.Columns(["bar", "baz", "foo"]), "(bar, baz, foo)")
def testSortsRawNamesWithoutEscape(self): self.assertGreater("`", "_") self.assertEqual(mysql_utils.Columns(["a", "a_hash"]), "(`a`, `a_hash`)")
def testSortsNames(self): self.assertEqual( mysql_utils.Columns(["bar", "foo", "baz"]), "(`bar`, `baz`, `foo`)")
def testMany(self): self.assertEqual( mysql_utils.Columns(["bar", "baz", "foo"]), "(`bar`, `baz`, `foo`)")
def testOne(self): self.assertEqual(mysql_utils.Columns(["foo"]), "(`foo`)")
def testEmpty(self): self.assertEqual(mysql_utils.Columns([]), "()")
def testDictUsesKeys(self): self.assertIn(mysql_utils.Columns({ "bar": 42, "baz": 42, "foo": 42 }), ["(bar, baz, foo)"])
def testSortsNames(self): self.assertEqual(mysql_utils.Columns(["bar", "foo", "baz"]), "(bar, baz, foo)")