Refactor into a proper class

Change-Id: If1d283b02ee72fae1709bba80755a342e684c320