Small Refactor

Yeah, it's dodgy as all getout, but at least it works!