Python codeql queries to get flows from usages of a class member/field

[yuval-piiano]

Hi,

I have this code :

import logging
from dataclasses import dataclass

logger = logging.getLogger(__name__)

class AddressData:
    address: str

ad = AddressData("my address")
logger.warning(ad.address)
intermedia_var = ad.first_name
logger.warning("first_name: %s", intermedia_var)

I would like to get all flows from the fields address in AddressData to the logger calls,

when I annotate AddressData with @DataClass the answe to #12313 helps :

import python
import semmle.python.dataflow.new.DataFlow
import semmle.python.ApiGraphs

class AddressData extends API::EntryPoint {
  AddressData() { this = "AddressData" }

  override DataFlow::LocalSourceNode getASource() {
    result.asExpr().(ClassExpr).getName() = "AddressData"
  }

  // Handles the decorator `@dataclass`
  API::Node getDataclass() {
    exists(API::CallNode decorate |
      // decoration happens when invoking the `@dataclass` decorator
      decorate = API::moduleImport("dataclasses").getMember("dataclass").getACall() and
      // the first argument to the decorator is the `AddressData` class
      decorate.getArg(0) = this.getASource() and
      // the decorated class is the result of the decorator
      result = decorate.getReturn()
    )
  }
}

from API::Node address, API::CallNode warning
where
  // take the data class, instantiate it and get the `address` field
  address = any(AddressData ad).getDataclass().getReturn().getMember("address") and
  // create a logger and call `warning` (you could add a list of interesting method names here)
  warning =
    API::moduleImport("logging").getMember("getLogger").getReturn().getMember("warning").getACall() and
  // require the `address` field to be and argument to the call to `warning`
  warning.getArg(_) = address.getAValueReachableFromSource()
select warning

mark the relevant code, but without @DataClass the query doesn't work:

from API::Node address, API::CallNode warning
where
  // take the data class, instantiate it and get the `address` field
  addresss = any(AddressData ad).getASource().getReturn().getMember("address")
  // create a logger and call `warning` (you could add a list of interesting method names here)
  warning =
    API::moduleImport("logging").getMember("getLogger").getReturn().getMember("warning").getACall() and
  // require the `address` field to be and argument to the call to `warning`
  warning.getArg(_) = address.getAValueReachableFromSource()
select warning

getting - "getReturn() cannot be resolved for type LocalSources::LocalSourceNode"
what's the correct way to do this?

[RasmusWL]

Hi @yuval-piiano. The code that @yoff provided in #12313 was slightly wrong, since you need to use getANode() and not getASource() on any(AddressData ad)