Embedding screenshots in Cucumber JVM

I was trying to improve our automated test suite and I thought that it could be useful to capture a screenshot of the browser whenever a test fails.
The current Cucumber JVM implementation highly simplifies this task, but the task is not achieved the way I thought, so this is the reason for this post.

Normally you would use a JUnit TestRule to augment all your test cases with a feature to take screenshots, but for Cucumber JVM it's much easier thankfully to the Execution Hooks:

 public class ScreenshotHook {  
     @Autowired  
     private WebDriver driver;
 
     @After  
     public void embedScreenshot(Scenario scenario) {  
         if (scenario.isFailed()) {  
             try {  
                 byte[] screenshot = ((TakesScreenshot) driver).getScreenshotAs(OutputType.BYTES);  
                 scenario.embed(screenshot, "image/png");  
             } catch (WebDriverException wde) {  
                 System.err.println(wde.getMessage());  
             } catch (ClassCastException cce) {  
                 cce.printStackTrace();  
             }  
         }  
     }  
 }  

As you can see I'm injecting the WebDriver instance (I'm using the Cucumber JVM Spring integration) and defining an embedScreenshot method which is annotated with cucumber.annotation.After: this is the important bit because methods annotated as such will be executed after each cucumber scenario.

For sake of completeness I want to tell you there is another hook cucumber.annotation.Before that can be used for things like logging into your application eliminating those repetitive log in steps.

UPDATE: Scenario was ScenarioResult in pre 1.0.0 versions of Cucumber JVM.