Refactoring some of the old source files to match the new classes