Refactoring/DRY up code