refactored code

Signed-off-by: Aubrey Howell <simplyaubs@gmail.com>