Upgrade main contract packege, audit fix and make token abstract

contracts/Migrations.sol

@@ -1,4 +1,4 @@
-pragma solidity >=0.4.21 <0.6.0;
+pragma solidity >=0.6.0 <0.8.0;
 
 
 contract Migrations {

contracts/StakingToken.sol

@@ -1,15 +1,16 @@
-pragma solidity ^0.5.0;
+pragma solidity >=0.6.0 <0.8.0;
+
+import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+import "@openzeppelin/contracts/math/SafeMath.sol";
+import "@openzeppelin/contracts/access/Ownable.sol";
 
-import "openzeppelin-solidity/contracts/token/ERC20/ERC20.sol";
-import "openzeppelin-solidity/contracts/math/SafeMath.sol";
-import "openzeppelin-solidity/contracts/ownership/Ownable.sol";
 
 /**
- * @title Staking Token (STK)
+ * @title Staking Token
  * @author Alberto Cuesta Canada
  * @notice Implements a basic ERC20 staking token with incentive distribution.
  */
-contract StakingToken is ERC20, Ownable {
+abstract contract StakingToken is ERC20, Ownable {
     using SafeMath for uint256;
 
     /**
@@ -29,13 +30,16 @@ contract StakingToken is ERC20, Ownable {
 
     /**
      * @notice The constructor for the Staking Token.
-     * @param _owner The address to receive all tokens on construction.
-     * @param _supply The amount of tokens to mint on construction.
+     * @param _name The name of the token.
+     * @param _symbol The symbol of the token.
      */
-    constructor(address _owner, uint256 _supply) 
+    constructor(
+        string memory _name,
+        string memory _symbol
+    )
         public
-    { 
-        _mint(_owner, _supply);
+        ERC20(_name, _symbol)
+    {
     }
 
     // ---------- STAKES ----------
@@ -48,7 +52,7 @@ contract StakingToken is ERC20, Ownable {
         public
     {
         _burn(msg.sender, _stake);
-        if(stakes[msg.sender] == 0) addStakeholder(msg.sender);
+        if (stakes[msg.sender] == 0) addStakeholder(msg.sender);
         stakes[msg.sender] = stakes[msg.sender].add(_stake);
     }
 
@@ -59,8 +63,10 @@ contract StakingToken is ERC20, Ownable {
     function removeStake(uint256 _stake)
         public
     {
+        require(stakes[msg.sender] >= _stake, "StakingToken: Insufficient amount of stakes");
+
         stakes[msg.sender] = stakes[msg.sender].sub(_stake);
-        if(stakes[msg.sender] == 0) removeStakeholder(msg.sender);
+        if (stakes[msg.sender] == 0) removeStakeholder(msg.sender);
         _mint(msg.sender, _stake);
     }
 
@@ -98,7 +104,7 @@ contract StakingToken is ERC20, Ownable {
     /**
      * @notice A method to check if an address is a stakeholder.
      * @param _address The address to verify.
-     * @return bool, uint256 Whether the address is a stakeholder, 
+     * @return bool, uint256 Whether the address is a stakeholder,
      * and if so its position in the stakeholders array.
      */
     function isStakeholder(address _address)
@@ -120,7 +126,7 @@ contract StakingToken is ERC20, Ownable {
         public
     {
         (bool _isStakeholder, ) = isStakeholder(_stakeholder);
-        if(!_isStakeholder) stakeholders.push(_stakeholder);
+        if (!_isStakeholder) stakeholders.push(_stakeholder);
     }
 
     /**
@@ -131,19 +137,19 @@ contract StakingToken is ERC20, Ownable {
         public
     {
         (bool _isStakeholder, uint256 s) = isStakeholder(_stakeholder);
-        if(_isStakeholder){
+        if (_isStakeholder){
             stakeholders[s] = stakeholders[stakeholders.length - 1];
             stakeholders.pop();
-        } 
+        }
     }
 
     // ---------- REWARDS ----------
-    
+
     /**
      * @notice A method to allow a stakeholder to check his rewards.
      * @param _stakeholder The stakeholder to check rewards for.
      */
-    function rewardOf(address _stakeholder) 
+    function rewardOf(address _stakeholder)
         public
         view
         returns(uint256)
@@ -167,7 +173,7 @@ contract StakingToken is ERC20, Ownable {
         return _totalRewards;
     }
 
-    /** 
+    /**
      * @notice A simple method that calculates the rewards for each stakeholder.
      * @param _stakeholder The stakeholder to calculate rewards for.
      */
@@ -182,7 +188,7 @@ contract StakingToken is ERC20, Ownable {
     /**
      * @notice A method to distribute rewards to all stakeholders.
      */
-    function distributeRewards() 
+    function distributeRewards()
         public
         onlyOwner
     {
@@ -196,7 +202,7 @@ contract StakingToken is ERC20, Ownable {
     /**
      * @notice A method to allow a stakeholder to withdraw his rewards.
      */
-    function withdrawReward() 
+    function withdrawReward()
         public
     {
         uint256 reward = rewards[msg.sender];

contracts/TestStakingToken.sol

@@ -0,0 +1,30 @@
+pragma solidity >=0.6.0 <0.8.0;
+
+import "./StakingToken.sol";
+
+
+/**
+ * @title Test Staking Token (TST)
+ * @author Dusan Perisic
+ * @notice Implementation of an abstract Staking Token.
+ */
+contract TestStakingToken is StakingToken {
+    /**
+     * @notice The constructor for the Staking Token.
+     * @param _name The name of the token.
+     * @param _symbol The symbol of the token.
+     * @param _owner The address to receive all tokens on construction.
+     * @param _supply The amount of tokens to mint on construction.
+     */
+    constructor(
+        string memory _name,
+        string memory _symbol,
+        address _owner,
+        uint256 _supply
+    )
+    public
+    StakingToken(_name, _symbol)
+    {
+        _mint(_owner, _supply);
+    }
+}

migrations/02_deploy_staking.js

@@ -1,9 +1,12 @@
 const BigNumber = require('bignumber.js');
-const StakingToken = artifacts.require('./StakingToken.sol');
 
-module.exports = function(deployer, network, accounts) {
+const TestStakingToken = artifacts.require('./TestStakingToken.sol');
+
+module.exports = (deployer, network, accounts) => {
     deployer.deploy(
-        StakingToken, 
+        TestStakingToken,
+        'Test Staking Token',
+        'TST',
         accounts[1],
         new BigNumber(10).pow(18).multipliedBy(525).toString(10),
     );